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

v4.5 dev update. (#3153)

Browse Source
This commit is contained in:
Junkai-Wu
2026-04-07 12:16:05 -04:00
committed by GitHub
parent 418d38a5de
commit a221da7ccf
265 changed files with 4913 additions and 1478 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/python/CuTeDSL/ampere/dynamic_smem_size.py
+44 -13
View File
@@ -26,9 +26,11 @@
# 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 cutlass.cute as cute
import cutlass
"""
Example of automatic shared memory size computation for configuring kernel launch
@@ -51,11 +53,15 @@ class SharedData:
@cute.kernel
def kernel():
def kernel_static():
"""
Example kernel that allocates shared memory.
The total allocation will be automatically calculated when smem=None.
"""
tidx, _, _ = cute.arch.block_idx()
if tidx == 0:
cute.printf("Running kernel_static")
allocator = cutlass.utils.SmemAllocator()
# Allocate various types of shared memory
@@ -68,6 +74,8 @@ def kernel():
byte_alignment=16,
swizzle=None,
)
cute.printf("Kernel launch smem size: {}", cute.arch.dynamic_smem_size())
return
@@ -79,7 +87,7 @@ def kernel_no_smem():
"""
tidx, _, _ = cute.arch.block_idx()
if tidx == 0:
cute.printf("Hello world")
cute.printf("Running kernel_no_smem")
return
@@ -89,26 +97,49 @@ if __name__ == "__main__":
print("Launching kernel with auto smem size. (launch config `smem=None`)")
# Compile the example
# Compile the static example
@cute.jit
def launch_kernel1():
k = kernel()
k.launch(
def launch_kernelno_smem():
kernel_no_smem().launch(
grid=(1, 1, 1),
block=(1, 1, 1),
)
print(f"Kernel recorded internal smem usage: {k.smem_usage()}")
# --------
print(f" > Run {kernel_no_smem.__name__}")
func = cute.compile(launch_kernelno_smem)
func()
cutlass.cuda.stream_sync(cutlass.cuda.default_stream())
@cute.jit
def launch_kernel2():
k = kernel_no_smem()
k.launch(
def launch_kernel_static():
kernel_static().launch(
grid=(1, 1, 1),
block=(1, 1, 1),
# smem=None
# auto infer launch kernel static smem usage
)
print(f"Kernel recorded internal smem usage: {k.smem_usage()}")
cute.compile(launch_kernel1)
cute.compile(launch_kernel2)
# --------
print(f" > Run {kernel_static.__name__} with sufficient smem")
func = cute.compile(launch_kernel_static)
func()
cutlass.cuda.stream_sync(cutlass.cuda.default_stream())
@cute.jit
def launch_kernel_static_insufficient():
kernel_static().launch(
grid=(1, 1, 1),
block=(1, 1, 1),
# launch kernel with static smem usage exceeds cfg
# show warning
smem=16,
)
# --------
print(f" > Run {kernel_static.__name__} with insufficient smem, show warning:")
func = cute.compile(launch_kernel_static_insufficient)
func()
cutlass.cuda.stream_sync(cutlass.cuda.default_stream())
print("PASS")
examples/python/CuTeDSL/ampere/hstu_attention.py
-1
View File
@@ -265,7 +265,6 @@ class HSTUAttentionForwardAmpere(object):
).launch(
grid=grid_dim,
block=[self._num_threads, 1, 1],
smem=SharedStorage.size_in_bytes(),
stream=stream,
)
examples/python/CuTeDSL/ampere/smem_allocator.py
+16 -3
View File
@@ -129,8 +129,8 @@ def kernel(
# ptr<i64, smem, align<128>>
# ptr<f32, smem, align<8>>
print(struct_in_smem.a.data_ptr())
print(struct_in_smem.b)
print(struct_in_smem.c.real)
print(struct_in_smem.b.ptr)
print(struct_in_smem.c.real.ptr)
# ptr<i8, smem, align<512>>
print(section_in_smem)
# ptr<i64, smem, align<64>>
@@ -138,6 +138,17 @@ def kernel(
# tensor<ptr<f16, smem, align<32>> o (16,4):(1,16)>
print(tensor_in_smem)
# assign struct member array element
cute.printf("struct_in_smem.a[0] = {}", struct_in_smem.a[0])
struct_in_smem.a[0] = 2
cute.printf("struct_in_smem.a[0] = {}", struct_in_smem.a[0])
# assign struct member scalar
cute.printf("struct_in_smem.b.ptr = {}", struct_in_smem.b.ptr)
cute.printf("struct_in_smem.b: value = {}", struct_in_smem.b.ptr.load())
struct_in_smem.b = 16
cute.printf("struct_in_smem.b: value = {}", struct_in_smem.b.ptr.load())
# fill MemRange tensor in struct and copy to dst
a_tensor = struct_in_smem.a.get_tensor(cute.make_layout((8, 4)))
a_tensor.fill(const_a)
@@ -169,7 +180,9 @@ def host(
):
# 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)
grid=(1, 1, 1),
block=(1, 1, 1),
# Automatically calculate the launch kernel shared memory usage when `smem=None`
)
examples/python/CuTeDSL/ampere/tensorop_gemm.py
+27 -15
View File
@@ -175,15 +175,6 @@ class TensorOpGemm:
(self.cta_tiler[0], self.cta_tiler[1]),
)
# Shared memory allocated for operations with A, B will be
# overwritten for operations on C. This is to improve performance
# by reducing the size of shared memory requested by each block
smem_size = max(
cute.size_in_bytes(mC.element_type, sC_layout),
cute.size_in_bytes(mA.element_type, sA_layout)
+ cute.size_in_bytes(mB.element_type, sB_layout),
)
# ///////////////////////////////////////////////////////////////////////////////
# Tiled copy:
# The majorness of tA/tB/tC follows the majorness of gA/gB/gC,
@@ -282,7 +273,6 @@ class TensorOpGemm:
).launch(
grid=rasterization_remap_grid_dim,
block=[self.num_threads, 1, 1],
smem=smem_size,
)
@cute.kernel
@@ -382,14 +372,36 @@ class TensorOpGemm:
# tAgA: (CPY, CPY_M, CPY_K, k) , tBgB: (CPY, CPY_N, CPY_K, k)
# tAsA: (CPY, CPY_M, CPY_K, PIPE) , tBsB: (CPY, CPY_N, CPY_K, PIPE)
# ///////////////////////////////////////////////////////////////////////////////
@cute.struct
class SharedStorageAB:
a: cute.struct.Align[
cute.struct.MemRange[mA.element_type, cute.cosize(sA_layout)],
16,
]
b: cute.struct.Align[
cute.struct.MemRange[mB.element_type, cute.cosize(sB_layout)],
16,
]
@cute.struct
class SharedStorageC:
c: cute.struct.Align[
cute.struct.MemRange[mC.element_type, cute.cosize(sC_layout)],
16,
]
# Shared memory buffer
smem = cutlass.utils.SmemAllocator()
sA = smem.allocate_tensor(mA.element_type, sA_layout, 16)
sB = smem.allocate_tensor(mB.element_type, sB_layout, 16)
sC = cute.make_tensor(
cute.recast_ptr(sA.iterator, dtype=self.c_dtype), sC_layout
# Shared memory allocated for operations with A, B will be
# overwritten for operations on C. This is to improve performance
# by reducing the size of shared memory requested by each block
storage = smem.allocate(
max(SharedStorageAB.size_in_bytes(), SharedStorageC.size_in_bytes()),
byte_alignment=16,
)
sA = SharedStorageAB(storage).a.get_tensor(sA_layout)
sB = SharedStorageAB(storage).b.get_tensor(sB_layout)
sC = SharedStorageC(storage).c.get_tensor(sC_layout)
thr_copy_A = tiled_copy_A.get_slice(tidx)
thr_copy_B = tiled_copy_B.get_slice(tidx)
examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py
+3 -7
View File
@@ -549,7 +549,7 @@ class BlockwiseGemmKernel:
cutlass.Int64, self.num_tile_stage * 2
]
epi_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 1 * 2]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sC: cute.struct.Align[
@@ -614,7 +614,6 @@ class BlockwiseGemmKernel:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
smem=self.shared_storage.size_in_bytes(),
stream=stream,
min_blocks_per_mp=1,
)
@@ -682,9 +681,6 @@ class BlockwiseGemmKernel:
smem = utils.SmemAllocator()
storage = smem.allocate(self.shared_storage)
tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar_ptr
tmem_holding_buf = storage.tmem_holding_buf
# Initialize mainloop ab_pipeline (barrier) and states
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
@@ -771,11 +767,11 @@ class BlockwiseGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py
+3 -7
View File
@@ -568,7 +568,7 @@ class BlockwiseContiguousGroupedGemmKernel:
cutlass.Int64, self.num_tile_stage * 2
]
epi_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 1 * 2]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sC: cute.struct.Align[
@@ -634,7 +634,6 @@ class BlockwiseContiguousGroupedGemmKernel:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
smem=self.shared_storage.size_in_bytes(),
stream=stream,
min_blocks_per_mp=1,
)
@@ -703,9 +702,6 @@ class BlockwiseContiguousGroupedGemmKernel:
smem = utils.SmemAllocator()
storage = smem.allocate(self.shared_storage)
tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar_ptr
tmem_holding_buf = storage.tmem_holding_buf
# Initialize mainloop ab_pipeline (barrier) and states
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
@@ -792,11 +788,11 @@ class BlockwiseContiguousGroupedGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py
+3 -7
View File
@@ -567,7 +567,7 @@ class BlockwiseMaskedGroupedGemmKernel:
cutlass.Int64, self.num_tile_stage * 2
]
epi_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 1 * 2]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sC: cute.struct.Align[
@@ -633,7 +633,6 @@ class BlockwiseMaskedGroupedGemmKernel:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
smem=self.shared_storage.size_in_bytes(),
stream=stream,
min_blocks_per_mp=1,
)
@@ -702,9 +701,6 @@ class BlockwiseMaskedGroupedGemmKernel:
smem = utils.SmemAllocator()
storage = smem.allocate(self.shared_storage)
tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar_ptr
tmem_holding_buf = storage.tmem_holding_buf
# Initialize mainloop ab_pipeline (barrier) and states
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
@@ -791,11 +787,11 @@ class BlockwiseMaskedGroupedGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py
+3 -3
View File
@@ -615,7 +615,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sC: cute.struct.Align[
@@ -794,11 +794,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent_amax.py
+3 -3
View File
@@ -551,7 +551,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sC: cute.struct.Align[
@@ -737,11 +737,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/dense_gemm.py
+3 -3
View File
@@ -517,7 +517,7 @@ class DenseGemmKernel:
acc_full_mbar_ptr: cute.struct.MemRange[
cutlass.Int64, self.num_acc_stage * 2
]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
smem = utils.SmemAllocator()
@@ -564,10 +564,10 @@ class DenseGemmKernel:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py
+3 -3
View File
@@ -493,7 +493,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
c_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage]
c_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sD: cute.struct.Align[
@@ -674,11 +674,11 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_ids[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
+3 -3
View File
@@ -593,7 +593,7 @@ class PersistentDenseGemmKernel:
acc_full_mbar_ptr: cute.struct.MemRange[
cutlass.Int64, self.num_acc_stage * 2
]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
smem = utils.SmemAllocator()
@@ -644,11 +644,11 @@ class PersistentDenseGemmKernel:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
allocator_warp_id=self.epilogue_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py
+3 -3
View File
@@ -614,7 +614,7 @@ class PersistentDenseGemmKernel:
acc_full_mbar_ptr: cute.struct.MemRange[
cutlass.Int64, self.num_acc_stage * 2
]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2]
clc_response: cute.struct.MemRange[cutlass.Int32, 4]
@@ -686,11 +686,11 @@ class PersistentDenseGemmKernel:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
allocator_warp_id=self.epilogue_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py
+3 -3
View File
@@ -514,7 +514,7 @@ class DenseGemmKernel:
acc_full_mbar_ptr: cute.struct.MemRange[
cutlass.Int64, self.num_acc_stage * 2
]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
smem = utils.SmemAllocator()
@@ -562,10 +562,10 @@ class DenseGemmKernel:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/epilogue/common_dense_gemm_efc.py
+3 -3
View File
@@ -573,7 +573,7 @@ class DenseGemmEFC:
# Barriers used by the supplemental load tensor pipeline.
c_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage]
c_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (MMA, MMA_M, MMA_K, STAGE)
sA: cute.struct.Align[
@@ -651,11 +651,11 @@ class DenseGemmEFC:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
allocator_warp_id=self.epilogue_warp_id[0],
is_two_cta=self.use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/fmha.py
+2 -2
View File
@@ -998,7 +998,7 @@ class BlackwellFusedMultiHeadAttentionForward:
# Alloc tmem buffer
tmem_alloc_cols = Int32(self.tmem_alloc_cols)
cute.arch.alloc_tmem(tmem_alloc_cols, storage.tmem_holding_buf)
cute.arch.alloc_tmem(tmem_alloc_cols, storage.tmem_holding_buf.ptr)
self.tmem_alloc_barrier.arrive_and_wait()
tile_sched = fmha_utils.create_fmha_static_tile_scheduler(
tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
@@ -1260,7 +1260,7 @@ class BlackwellFusedMultiHeadAttentionForward:
tmem_ptr = cute.arch.retrieve_tmem_ptr(
Float32,
alignment=16,
ptr_to_buffer_holding_addr=storage.tmem_holding_buf,
ptr_to_buffer_holding_addr=storage.tmem_holding_buf.ptr,
)
cute.arch.dealloc_tmem(tmem_ptr, tmem_alloc_cols)
examples/python/CuTeDSL/blackwell/fmha_bwd.py
+1 -2
View File
@@ -708,7 +708,6 @@ class BlackwellFusedMultiHeadAttentionBackward:
grid=bwd_grid,
block=[self.threads_per_cta, 1, 1],
cluster=[1, 1, 1],
smem=self.shared_storage.size_in_bytes(),
stream=stream,
min_blocks_per_mp=1,
)
@@ -913,7 +912,7 @@ class BlackwellFusedMultiHeadAttentionBackward:
)
sLSE = storage.sLSE.get_tensor(LSE_smem_layout)
sSum_OdO = storage.sSum_OdO.get_tensor(sum_OdO_smem_layout)
tmem_holding_buf = storage.tmem_holding_buf
tmem_holding_buf = storage.tmem_holding_buf.ptr
sQT_ptr = cute.recast_ptr(sQ.iterator, QT_smem_layout_staged.inner)
sQT = cute.make_tensor(sQT_ptr, QT_smem_layout_staged.outer)
examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py
+6 -7
View File
@@ -567,7 +567,7 @@ class Sm100GroupedBlockScaledGemmKernel:
ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sC: cute.struct.Align[
@@ -641,7 +641,6 @@ class Sm100GroupedBlockScaledGemmKernel:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
smem=self.shared_storage.size_in_bytes(),
stream=stream,
min_blocks_per_mp=1,
)
@@ -737,8 +736,8 @@ class Sm100GroupedBlockScaledGemmKernel:
+ Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8
)
tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar_ptr
tmem_holding_buf = storage.tmem_holding_buf
tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar.ptr
tmem_holding_buf_ptr = storage.tmem_holding_buf.ptr
# Initialize mainloop ab_pipeline (barrier) and states
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
@@ -1249,7 +1248,7 @@ class Sm100GroupedBlockScaledGemmKernel:
acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
self.acc_dtype,
alignment=16,
ptr_to_buffer_holding_addr=tmem_holding_buf,
ptr_to_buffer_holding_addr=tmem_holding_buf_ptr,
)
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
@@ -1446,7 +1445,7 @@ class Sm100GroupedBlockScaledGemmKernel:
if warp_idx == self.epilog_warp_id[0]:
cute.arch.alloc_tmem(
self.num_tmem_alloc_cols,
tmem_holding_buf,
tmem_holding_buf_ptr,
is_two_cta=use_2cta_instrs,
)
@@ -1461,7 +1460,7 @@ class Sm100GroupedBlockScaledGemmKernel:
acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
self.acc_dtype,
alignment=16,
ptr_to_buffer_holding_addr=tmem_holding_buf,
ptr_to_buffer_holding_addr=tmem_holding_buf_ptr,
)
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
examples/python/CuTeDSL/blackwell/grouped_gemm.py
+3 -3
View File
@@ -425,7 +425,7 @@ class GroupedGemmKernel:
ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sC: cute.struct.Align[
@@ -590,11 +590,11 @@ class GroupedGemmKernel:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py
+1 -1
View File
@@ -820,7 +820,7 @@ class SSDKernel:
num_threads=self.threads_per_cta,
)
tmem = utils.TmemAllocator(
smem_storage.tmem_holding_buf,
smem_storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
)
examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py
-2
View File
@@ -603,8 +603,6 @@ class MixedInputFusedMultiHeadAttentionDecode:
p_pipeline_ptr = smem.allocate_array(Int64, self.sp_stages * 2)
o_pipeline_ptr = smem.allocate_array(Int64, self.o_stages * 2)
assert smem._allocated_bytes <= self.mbarrier_reserved_bytes
# Declare named barriers
softmax_nbar_id = 1
mma_kq_nbar_id = 2
examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d256.py
+3 -3
View File
@@ -403,7 +403,7 @@ class MixedInputFusedMultiHeadAttentionPrefillD256:
s_corr_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2]
sum_mbar_ptr: cute.struct.MemRange[Int64, 2]
mma_o_mbar_ptr: cute.struct.MemRange[Int64, self.pv_acc_stage * 2]
tmem_dealloc_mbar_ptr: Int64
tmem_dealloc_mbar: Int64
tmem_holding_buf: Int32
self.shared_storage = SharedStorage
@@ -654,11 +654,11 @@ class MixedInputFusedMultiHeadAttentionPrefillD256:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
allocator_warp_id=self.correction_warp_ids[0],
is_two_cta=True,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True)
examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill_d512.py
+3 -3
View File
@@ -390,7 +390,7 @@ class MixedInputFusedMultiHeadAttentionPrefillD512:
p_mma_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2]
mma_o_mbar_ptr: cute.struct.MemRange[Int64, self.pv_acc_stage * 2]
swap_mbar_ptr: cute.struct.MemRange[Int64, self.swap_stage * 2]
tmem_dealloc_mbar_ptr: Int64
tmem_dealloc_mbar: Int64
tmem_holding_buf: Int32
self.shared_storage = SharedStorage
@@ -627,11 +627,11 @@ class MixedInputFusedMultiHeadAttentionPrefillD512:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
allocator_warp_id=self.softmax_warp_ids[0],
is_two_cta=True,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True)
examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm.py
+3 -3
View File
@@ -624,7 +624,7 @@ class GroupedMixedInputGemmKernel:
tile_info_empty_mbar_ptr: cute.struct.MemRange[
cutlass.Int64, self.num_tile_info_stage
]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
self.shared_storage = SharedStorage
@@ -824,11 +824,11 @@ class GroupedMixedInputGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_ptr_sync_barrier,
allocator_warp_id=self.epilog_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/mixed_input_gemm/grouped_mixed_input_gemm_acc_scale.py
+3 -3
View File
@@ -518,7 +518,7 @@ class GroupedMixedInputGemmAccScaleKernel:
tile_info_empty_mbar_ptr: cute.struct.MemRange[
cutlass.Int64, self.num_tile_info_stage
]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
self.shared_storage = SharedStorage
@@ -708,11 +708,11 @@ class GroupedMixedInputGemmAccScaleKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_ptr_sync_barrier,
allocator_warp_id=self.epilog_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/mixed_input_gemm/mixed_input_gemm.py
+3 -3
View File
@@ -618,7 +618,7 @@ class MixedInputGemmKernel:
]
acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# Tensor buffers
# (EPI_TILE_M, EPI_TILE_N, STAGE)
@@ -820,11 +820,11 @@ class MixedInputGemmKernel:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_ptr_sync_barrier,
allocator_warp_id=self.epilog_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/blackwell/mla/mla_decode_fp16.py
+3 -5
View File
@@ -568,7 +568,7 @@ class BlackwellMultiHeadLatentAttentionForwardFP16:
cutlass.Int64, self.load_pt_stage * 2
]
# Tmem dealloc cluster barrier
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
# Tmem holding buffer
tmem_holding_buf: cutlass.Int32
@@ -641,7 +641,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP16:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=self.cluster_shape_mnk,
smem=SplitKVKernelSharedStorage.size_in_bytes(),
stream=stream,
min_blocks_per_mp=1,
)
@@ -657,7 +656,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP16:
).launch(
grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]),
block=[self.threads_per_warp * self.num_compute_warps, 1, 1],
smem=MAX_SPLITS * self.acc_dtype.width // 8,
stream=stream,
min_blocks_per_mp=1,
)
@@ -838,11 +836,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP16:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_ptr_sync_bar,
allocator_warp_id=self.mma_warp_id,
is_two_cta=self.use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
load_q_pipeline = self.make_and_init_load_qkv_pipeline(
examples/python/CuTeDSL/blackwell/mla/mla_decode_fp8.py
+3 -5
View File
@@ -661,7 +661,7 @@ class BlackwellMultiHeadLatentAttentionForwardFP8:
]
# Tmem dealloc cluster barrier
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
# Tmem holding buffer
tmem_holding_buf: cutlass.Int32
@@ -707,7 +707,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP8:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=self.cluster_shape_mnk,
smem=SplitKVKernelSharedStorage.size_in_bytes(),
stream=stream,
min_blocks_per_mp=1,
)
@@ -723,7 +722,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP8:
).launch(
grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]),
block=[self.threads_per_warp * self.num_compute_warps, 1, 1],
smem=MAX_SPLITS * self.acc_dtype.width // 8,
stream=stream,
min_blocks_per_mp=1,
)
@@ -904,11 +902,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP8:
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=self.tmem_ptr_sync_bar,
allocator_warp_id=self.mma_warp_id,
is_two_cta=self.use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
load_q_pipeline = self.make_and_init_load_qkv_pipeline(
examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py
+9 -9
View File
@@ -43,9 +43,9 @@ from cutlass.cute.runtime import from_dlpack
from dataclasses import dataclass, field
"""
This example provides an experimental implementation of the SM103 batched 3xFP4 blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases.
This example provides an experimental implementation of the SM103 batched FP4 Ultra blockscaled GEMM kernel, please note that the APIs and implementation details related to this kernel may change in future releases.
A high-performance persistent batched 3xFP4 blockscaled GEMM example for the NVIDIA Blackwell SM103 architecture
A high-performance persistent batched FP4 Ultra blockscaled GEMM example for the NVIDIA Blackwell SM103 architecture
using CUTE DSL.
- Matrix A is MxKxL, L is batch dimension, A can only be row-major("K") for MXF4/NVF4 input type
- Matrix B is NxKxL, L is batch dimension, B can only be row-major("K") for MXF4/NVF4 input type
@@ -166,7 +166,7 @@ class Sm103BlockScaledPersistentDenseGemmKernel:
cluster_shape_mn: Tuple[int, int],
use_tma_store: bool,
):
"""Initializes the configuration for a Blackwell SM103 3xFP4 GEMM kernel.
"""Initializes the configuration for a Blackwell SM103 FP4 Ultra GEMM kernel.
This configuration includes several key aspects:
@@ -603,7 +603,7 @@ class Sm103BlockScaledPersistentDenseGemmKernel:
sf_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_sf_stage]
acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (MMA, MMA_M, MMA_K, STAGE)
sA: cute.struct.Align[
@@ -800,11 +800,11 @@ class Sm103BlockScaledPersistentDenseGemmKernel:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
allocator_warp_id=self.epilogue_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
@@ -1810,7 +1810,7 @@ class Sm103BlockScaledPersistentDenseGemmKernel:
mma_tiler_mn: Tuple[int, int],
a_source: tcgen05.OperandSource = tcgen05.OperandSource.SMEM,
) -> cute.TiledMma:
"""Create a blockscaled trivial tiled MMA for SM103 (3xFP4), K fixed to 96.
"""Create a blockscaled trivial tiled MMA for SM103 (FP4 Ultra), K fixed to 96.
Returns a tcgen05 MMA configured for the given (M, N) tiler and CTA group.
@@ -2653,7 +2653,7 @@ def run(
:return: Execution time of the GEMM kernel
:rtype: float
"""
print(f"Running Sm103 Persistent 3xfp4 Dense BlockScaled GEMM test with:")
print(f"Running Sm103 Persistent FP4 Ultra Dense BlockScaled GEMM test with:")
print(f"mnkl: {mnkl}")
print(f"AB dtype: {ab_dtype}, SF dtype: {sf_dtype}, SF Vec size: {sf_vec_size}")
print(f"C dtype: {c_dtype}")
@@ -2954,7 +2954,7 @@ if __name__ == "__main__":
)
parser = argparse.ArgumentParser(
description="Example of Sm103 3xfp4 Dense Persistent BlockScaled GEMM."
description="Example of Sm103 FP4 Ultra Dense Persistent BlockScaled GEMM."
)
parser.add_argument(
examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_0.py
+49 -64
View File
@@ -8,13 +8,10 @@
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
import argparse
from typing import Tuple, Type, Callable
from functools import partial, lru_cache
from typing import Tuple
import cutlass
from cutlass import Numeric
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.pipeline as pipeline
@@ -69,7 +66,6 @@ def kernel(
a_smem_layout: cute.ComposedLayout,
b_smem_layout: cute.ComposedLayout,
):
# Current thread/warp/block coordinates
tidx, _, _ = cute.arch.thread_idx()
warp_idx = cute.arch.warp_idx()
@@ -103,7 +99,7 @@ def kernel(
num_threads=threads_per_cta,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
)
num_tmem_cols = 512
@@ -143,15 +139,15 @@ def kernel(
# (bM, bN)
gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None))
thr_mma = tiled_mma.get_slice(0)
# (MMA, MMA_M, MMA_K, RestK)
# (MMA, MMA_M, MMA_K)
tCgA = thr_mma.partition_A(gA)
# (MMA, MMA_N, MMA_K, RestK)
# (MMA, MMA_N, MMA_K)
tCgB = thr_mma.partition_B(gB)
# (MMA, MMA_M, MMA_N)
tCgC = thr_mma.partition_C(gC)
# (MMA, MMA_M, MMA_K, STAGE)
# (MMA, MMA_M, MMA_K)
tCrA = tiled_mma.make_fragment_A(sA)
# (MMA, MMA_N, MMA_K, STAGE)
# (MMA, MMA_N, MMA_K)
tCrB = tiled_mma.make_fragment_B(sB)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
@@ -199,14 +195,14 @@ def kernel(
tmem_thr_copy = tmem_tiled_copy.get_slice(tidx)
# (TmemCpy,NumTmemCpy,NumTiles)
tCtC = tmem_thr_copy.partition_S(tCtAcc_epi)
tDtC = tmem_thr_copy.partition_S(tCtAcc_epi)
# (TmemCpy,NumTmemCpy,NumTiles)
tCgC = tmem_thr_copy.partition_D(gC_epi)
tDgC = tmem_thr_copy.partition_D(gC_epi)
# (TmemCpy,NumTmemCpy)
tCrAcc = cute.make_rmem_tensor(tCgC[None, None, 0].shape, acc_dtype)
tCrAcc = cute.make_rmem_tensor(tDgC[None, None, 0].shape, acc_dtype)
# (TmemCpy,NumTmemCpy)
tCrC = cute.make_rmem_tensor(tCgC[None, None, 0].shape, io_dtype)
tCrC = cute.make_rmem_tensor(tDgC[None, None, 0].shape, io_dtype)
#
# 2. Main loop
@@ -233,8 +229,6 @@ def kernel(
# Execute one K-block worth of MMA instructions
ab_full = ab_consumer.wait_and_advance()
# tCtAcc += tCrA * tCrB
num_k_blocks = cute.size(tCrA, mode=[2])
for k_block_idx in cutlass.range_constexpr(num_k_blocks):
k_block_coord = (None, None, k_block_idx, ab_full.index)
@@ -265,10 +259,10 @@ def kernel(
# TMEM -> RMEM -> GEMM
# Sub-tiling for better instruction-level parallelism
for i in cutlass.range(cute.size(tCtC, mode=[2])):
cute.copy(tmem_tiled_copy, tCtC[None, None, i], tCrAcc)
for i in cutlass.range(cute.size(tDtC, mode=[2])):
cute.copy(tmem_tiled_copy, tDtC[None, None, i], tCrAcc)
tCrC.store(tCrAcc.load().to(io_dtype))
cute.autovec_copy(tCrC, tCgC[None, None, i])
cute.autovec_copy(tCrC, tDgC[None, None, i])
acc_full.release()
# Deallocate TMEM
@@ -350,44 +344,10 @@ def host_function(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
)
@lru_cache(maxsize=1)
def prepare_run(
callable: Callable,
m: int,
n: int,
k: int,
a_dtype: Type[Numeric],
b_dtype: Type[Numeric],
c_dtype: Type[Numeric],
) -> tuple[Callable, tuple]:
import cutlass.torch as cutlass_torch
a, b, c = cutlass_torch.prepare_tensors_for_gemm(
(m, n, k), a_dtype, b_dtype, c_dtype
)
a_ = (
from_dlpack(a, assumed_align=32)
.mark_layout_dynamic(leading_dim=1)
.mark_compact_shape_dynamic(mode=1, divisibility=k)
)
b_ = (
from_dlpack(b, assumed_align=32)
.mark_layout_dynamic(leading_dim=1)
.mark_compact_shape_dynamic(mode=1, divisibility=k)
)
c_ = (
from_dlpack(c, assumed_align=32)
.mark_layout_dynamic(leading_dim=1)
.mark_compact_shape_dynamic(mode=1, divisibility=n)
)
compiled_fn = cute.compile(callable, a_, b_, c_, options="--generate-line-info")
return partial(compiled_fn, a_, b_, c_), (a, b, c)
def run_dense_gemm(
mnk: Tuple[int, int, int],
tolerance: float,
) -> None:
):
global torch, cutlass_torch
import torch
import cutlass.torch as cutlass_torch
@@ -402,23 +362,48 @@ def run_dense_gemm(
m, n, k = mnk
torch.manual_seed(1111)
run_fn, (a, b, c) = prepare_run(
host_function, m, n, k, io_dtype, io_dtype, io_dtype
# Make K-major tensors (torch tensors are row-major)
def make_tensors(mn, k, dtype):
shape = (mn, k)
return (
torch.empty(*shape, dtype=torch.int32)
.random_(-2, 2)
.to(dtype=dtype, device="cuda")
)
a = make_tensors(m, k, cutlass_torch.dtype(io_dtype))
b = make_tensors(n, k, cutlass_torch.dtype(io_dtype))
c = make_tensors(m, n, cutlass_torch.dtype(io_dtype))
a_tensor = (
from_dlpack(a, assumed_align=32)
.mark_layout_dynamic(leading_dim=1)
.mark_compact_shape_dynamic(mode=1, divisibility=k)
)
b_tensor = (
from_dlpack(b, assumed_align=32)
.mark_layout_dynamic(leading_dim=1)
.mark_compact_shape_dynamic(mode=1, divisibility=k)
)
c_tensor = (
from_dlpack(c, assumed_align=32)
.mark_layout_dynamic(leading_dim=1)
.mark_compact_shape_dynamic(mode=1, divisibility=n)
)
# Entry point to the host JIT function
run_fn()
host_function(a_tensor, b_tensor, c_tensor, no_cache=True)
# Compute reference result and verify
ref = torch.einsum("mk,nk->mn", a.to(torch.float32), b.to(torch.float32))
ref = (torch.einsum("mk,nk->mn", a.to(torch.float32), b.to(torch.float32))).cpu()
torch.testing.assert_close(
c, ref.to(cutlass_torch.dtype(io_dtype)), atol=tolerance, rtol=1e-05
c.cpu(), ref.to(cutlass_torch.dtype(io_dtype)), atol=tolerance, rtol=1e-05
)
if __name__ == "__main__":
def parse_comma_separated_ints(s: str) -> list[int]:
def parse_comma_separated_ints(s: str):
try:
return [int(x.strip()) for x in s.split(",")]
except ValueError:
@@ -443,14 +428,14 @@ if __name__ == "__main__":
parser.add_argument(
"--tolerance", type=float, default=1e-01, help="Tolerance for validation"
)
args = parser.parse_args()
if len(args.mnk) != 3:
parser.error("--mnk must contain exactly 3 values")
if args.mnk[0] % mma_tiler_mnk[0] != 0 or args.mnk[1] % mma_tiler_mnk[1] != 0:
parser.error("m n must be divisible by mma_tiler_mn")
run_dense_gemm(args.mnk, args.tolerance)
run_dense_gemm(
args.mnk,
args.tolerance,
)
print("PASS")
examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_1.py
+22 -32
View File
@@ -65,8 +65,7 @@ Constraints for this example:
io_dtype = cutlass.Float16
acc_dtype = cutlass.Float32
use_2cta_instrs = True
cluster_shape_mnk = (2, 1, 1) if use_2cta_instrs else (1, 1, 1)
cluster_shape_mnk = (2, 1, 1)
mma_inst_shape_mnk = (256, 256, 16)
mma_tiler_mnk = (256, 256, 64)
threads_per_cta = 128
@@ -96,7 +95,6 @@ def kernel(
b_smem_layout: cute.ComposedLayout,
cta_layout_vmnk: cute.Layout,
):
# Current thread/warp/block coordinates
tidx, _, _ = cute.arch.thread_idx()
warp_idx = cute.arch.warp_idx()
@@ -174,15 +172,15 @@ def kernel(
# (bM, bN)
gC = cute.local_tile(mC_mnl, mma_tiler_mnk, mma_coord_mnk, proj=(1, 1, None))
thr_mma = tiled_mma.get_slice(mma_coord_vmnk[0])
# (MMA, MMA_M, MMA_K, RestK)
# (MMA, MMA_M, MMA_K)
tCgA = thr_mma.partition_A(gA)
# (MMA, MMA_N, MMA_K, RestK)
# (MMA, MMA_N, MMA_K)
tCgB = thr_mma.partition_B(gB)
# (MMA, MMA_M, MMA_N)
tCgC = thr_mma.partition_C(gC)
# (MMA, MMA_M, MMA_K, STAGE)
# (MMA, MMA_M, MMA_K)
tCrA = tiled_mma.make_fragment_A(sA)
# (MMA, MMA_N, MMA_K, STAGE)
# (MMA, MMA_N, MMA_K)
tCrB = tiled_mma.make_fragment_B(sB)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
@@ -217,10 +215,10 @@ def kernel(
num_threads=threads_per_cta,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
is_two_cta=cute.size(cta_layout_vmnk, mode=[0]) > 1,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
num_tmem_cols = 512
tmem.allocate(num_tmem_cols)
@@ -232,7 +230,7 @@ def kernel(
# Swap the pointer in tCtAcc
tCtAcc = cute.make_tensor(tmem_ptr, tCtAcc.layout)
subtile_cnt = 1 if mma_tiler_mnk[0] == 64 else 4
subtile_cnt = 4
# (EpiTile)
epi_tiler = (
(cute.size(tCtAcc, mode=[0, 0]), cute.size(tCtAcc, mode=[0, 1]) // subtile_cnt),
@@ -244,24 +242,21 @@ def kernel(
# Every thread loads 64 x fp32
tmem_atom = cute.make_copy_atom(
tcgen05.Ld16x256bOp(tcgen05.Repetition.x8)
if mma_tiler_mnk[0] == 64
else tcgen05.Ld32x32bOp(tcgen05.Repetition.x64),
tcgen05.Ld32x32bOp(tcgen05.Repetition.x64),
cutlass.Float32,
)
tmem_tiled_copy = tcgen05.make_tmem_copy(tmem_atom, tCtAcc_epi[None, 0])
tmem_thr_copy = tmem_tiled_copy.get_slice(tidx)
# (TmemCpy,NumTmemCpy,NumTiles)
tCtC = tmem_thr_copy.partition_S(tCtAcc_epi)
tDtC = tmem_thr_copy.partition_S(tCtAcc_epi)
# (TmemCpy,NumTmemCpy,NumTiles)
tCgC = tmem_thr_copy.partition_D(gC_epi)
tDgC = tmem_thr_copy.partition_D(gC_epi)
# (TmemCpy,NumTmemCpy)
tCrAcc = cute.make_rmem_tensor(tCgC[None, None, 0].shape, acc_dtype)
tCrAcc = cute.make_rmem_tensor(tDgC[None, None, 0].shape, acc_dtype)
# (TmemCpy,NumTmemCpy)
tCrC = cute.make_rmem_tensor(tCgC[None, None, 0].shape, io_dtype)
tCrC = cute.make_rmem_tensor(tDgC[None, None, 0].shape, io_dtype)
#
# 2. Main loop
@@ -271,8 +266,8 @@ def kernel(
if warp_idx == 0:
# Wait for a empty accumulator buffer
if is_leader_cta:
acc_producer.acquire()
for k_tile in cutlass.range(num_k_tiles, prefetch_stages=ab_stages - 2):
acc_producer.acquire_and_advance()
for _ in cutlass.range(num_k_tiles, prefetch_stages=ab_stages - 2):
# Issue TMA loads
ab_empty = ab_producer.acquire_and_advance()
cute.copy(
@@ -310,7 +305,6 @@ def kernel(
# Signal that the accumulator is fully computed
if is_leader_cta:
acc_producer.commit()
acc_producer.advance()
#
# 3. Epilogue
@@ -321,13 +315,12 @@ def kernel(
# Wait for the accumulator buffer to be full
acc_full = acc_consumer.wait_and_advance()
# TMEM -> RMEM -> GEMM
# Sub-tiling for better instruction-level parallelism
for i in cutlass.range(cute.size(tCtC, mode=[2])):
cute.copy(tmem_tiled_copy, tCtC[None, None, i], tCrAcc)
for i in cutlass.range(cute.size(tDtC, mode=[2])):
cute.copy(tmem_tiled_copy, tDtC[None, None, i], tCrAcc)
tCrC.store(tCrAcc.load().to(io_dtype))
cute.autovec_copy(tCrC, tCgC[None, None, i])
cute.autovec_copy(tCrC, tDgC[None, None, i])
acc_full.release()
# Ensure used buffers are properly synchronized before producer exit.
@@ -353,7 +346,7 @@ def host_function(
io_dtype,
acc_dtype,
mma_inst_shape_mnk,
tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE,
tcgen05.CtaGroup.TWO,
tcgen05.OperandSource.SMEM,
tcgen05.OperandMajorMode.K,
tcgen05.OperandMajorMode.K,
@@ -381,16 +374,14 @@ def host_function(
cta_layout_vmnk = cute.tiled_divide(cta_layout_mnk, (tiled_mma.thr_id,))
# Construct TMA load atoms
op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SMulticastOp(
tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE
)
op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.TWO)
a_tma_atom, a_tma_tensor = cute.nvgpu.make_tiled_tma_atom_A(
op,
a,
a_smem_layout_one_stage,
mma_tiler_mnk,
tiled_mma,
cta_layout_vmnk.shape,
cta_layout_vmnk.shape, # take the layout and extract the shape internally
)
b_tma_atom, b_tma_tensor = cute.nvgpu.make_tiled_tma_atom_B(
op,
@@ -403,8 +394,7 @@ def host_function(
grid_shape = cute.round_up(
cute.ceil_div(
(*c.layout.shape, 1),
(mma_tiler_mnk[0] // (2 if use_2cta_instrs else 1), *mma_tiler_mnk[1:]),
(*c.layout.shape, 1), (mma_tiler_mnk[0] // 2, *mma_tiler_mnk[1:])
),
cluster_shape_mnk,
)
examples/python/CuTeDSL/blackwell/tutorial_gemm/fp16_gemm_4.py
+1 -1
View File
@@ -981,7 +981,7 @@ def run_dense_gemm(
import cutlass.torch as cutlass_torch
print("===================================================================")
print("Running Blackwell fp16 GEMM example 4 (with MIX CGA support):")
print("Running Blackwell fp16 GEMM example 4 (with MIX cluster size support):")
print(f" mnk: {mnk}")
print(f" tolerance: {tolerance}")
print(f" Preferred cluster shape: {preferred_cluster_shape_mnk}")
examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_0.py
+1 -1
View File
@@ -500,7 +500,7 @@ class Sm100BlockScaledDenseGemmKernel:
num_threads=self.threads_per_cta,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
)
tmem.allocate(self.num_tmem_alloc_cols)
examples/python/CuTeDSL/blackwell/tutorial_gemm/nvfp4_gemm_1.py
+3 -3
View File
@@ -436,7 +436,7 @@ class Sm100BlockScaledDenseGemmKernel:
class SharedStorage:
ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2]
acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
smem = utils.SmemAllocator()
@@ -638,10 +638,10 @@ class Sm100BlockScaledDenseGemmKernel:
num_threads=self.threads_per_cta,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
is_two_cta=cute.size(cta_layout_vmnk, mode=[0]) > 1,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
tmem.allocate(self.num_tmem_alloc_cols)
tmem.wait_for_alloc()
examples/python/CuTeDSL/experimental/blackwell/dense_gemm_cute_pipeline.py
+3 -3
View File
@@ -581,7 +581,7 @@ class PersistentDenseGemmKernel:
acc_full_mbar_ptr: cute.struct.MemRange[
cutlass.Int64, self.num_acc_stage * 2
]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_dealloc_mbar: cutlass.Int64
tmem_holding_buf: cutlass.Int32
smem = utils.SmemAllocator()
@@ -632,11 +632,11 @@ class PersistentDenseGemmKernel:
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
storage.tmem_holding_buf.ptr,
barrier_for_retrieve=tmem_alloc_barrier,
allocator_warp_id=self.epilogue_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr,
)
# Cluster arrive after barrier init
examples/python/CuTeDSL/helpers/fmha_helpers.py
+1 -1
View File
@@ -3,7 +3,7 @@
#
# 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
# https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
examples/python/CuTeDSL/hopper/fmha.py
-1
View File
@@ -527,7 +527,6 @@ class HopperFusedMultiHeadAttentionForward:
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=self.cluster_shape_mnk,
smem=self.shared_storage.size_in_bytes(),
stream=stream,
min_blocks_per_mp=1,
)
examples/python/CuTeDSL/hopper/grouped_gemm.py
+1 -1
View File
@@ -107,7 +107,7 @@ Constraints (same as dense_gemm_persistent.py plus):
* Cluster shape M/N: power of 2, total <= 4
* Contiguous dim must be 16-byte aligned
Debug environment knobs:
Debug environment options:
* `GROUPED_GEMM_FORCE_CUTE_COPY=1`
Disable the non-mcast NVVM TMA load path and always use `cute.copy`.
"""
examples/python/CuTeDSL/jax/cutlass_call_basic.py
+12 -18
View File
@@ -30,7 +30,6 @@ from functools import partial
import jax
import jax.numpy as jnp
import cutlass
import cutlass.cute as cute
import cutlass.jax as cjax
import cuda.bindings.driver as cuda
@@ -140,12 +139,12 @@ if __name__ == "__main__":
def run_cutlass_kernel(a, b, x, y):
call = cjax.cutlass_call(
launch_jax_wrapper,
# Jax requires output shapes/dtype information for each output
# Describe the shape and dtype of each output buffer.
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
# Static jit arguments are passed via additional keyword arguments
# Static jit arguments are passed via additional keyword arguments.
x=x,
y=y,
)
@@ -165,12 +164,11 @@ if __name__ == "__main__":
# to the kernel. Alternatively you can wrap using another separate cute.jit
# function.
lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
# Jax requires output shapes/dtype information for each output
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
# Static jit arguments are passed via additional keyword arguments
# Static jit arguments are passed via additional keyword arguments.
x=x,
y=y,
)
@@ -191,11 +189,12 @@ if __name__ == "__main__":
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
# By default cutlass_call will treat all tensors as dynamic shape.
# By default cutlass_call treats all tensors as dynamic shape.
# Dynamic shapes are often expected for kernels so this default ensures
# the broadest support. If you know that a kernel can accept fully static
# tensors then you can enable this flag to pass all tensors shapes and
# layouts known at compile time.
# tensors then you can enable this flag to compile all tensor shapes and
# layouts as constexpr values known at compile time.
# Individual tensors may opt out via .mark_layout_dynamic().
use_static_tensors=True,
x=x,
y=y,
@@ -209,19 +208,15 @@ if __name__ == "__main__":
@partial(jax.jit, static_argnums=[2, 3])
def run_cutlass_kernel_with_modes(a, b, x, y):
# input_spec and output_spec accept TensorSpec values to attach layout
# metadata to tensors. mode remaps the logical dimension order seen by
# the kernel. static=True compiles that tensor's layout as constexpr.
call = cjax.cutlass_call(
lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
# The modes of the layout for each tensor can be specified using the
# TensorSpec. By default modes will align with the physical layout
# but can be mapped to specific index position. If None is passed
# then the default mode is assumed for that tensor.
#
# Individual static/dynamic settings may also be applied. For example
# a specific tensor can be marked to have static shape.
input_spec=(
cjax.TensorSpec(mode=(1, 0, 2), static=True),
cjax.TensorSpec(mode=(3, 1, 2, 0)),
@@ -245,9 +240,8 @@ if __name__ == "__main__":
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, b.dtype),
),
# Can specify the input tensors that are aliasing outputs of this call.
# To avoid allocating separate output buffers. This is useful for kernels
# that update a tensor.
# Map input indices to output indices so XLA can reuse the input
# buffers for the outputs, avoiding extra allocations.
input_output_aliases={0: 0, 1: 1},
x=x,
y=y,
examples/python/CuTeDSL/jax/cutlass_call_export.py
+114 -81
View File
@@ -26,45 +26,45 @@
# 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 pytest
from functools import partial
import argparse
import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
import jax
import jax.numpy as jnp
from jax import export
from cutlass.jax import cutlass_call, get_export_disabled_safety_checks
from cutlass.jax.testing import create_tensor
"""
Examples of using jax.export APIs with functions using cutlass_call.
This example demonstrates the use of jax.export with CuTe DSL kernel. It assumes
familiarity with CuTe DSL concepts such as layouts and dynamic shapes as well as
Jax's exporting and serialization features:
This example demonstrates three export modes:
1. Concrete shapes -- shapes are fixed constants baked into the export.
2. Unconstrained symbolic shapes ("a, b")
3. Constrained symbolic shapes ("32*M, 16*N")
The JAX function being exported is the same in all three cases; only the
shape specification passed to jax.export differs.
It assumes familiarity with CuTe DSL concepts such as layouts and dynamic shapes
as well as JAX's exporting and serialization features:
https://docs.jax.dev/en/latest/export/index.html#export
To run this example:
.. code-block:: bash
# Run with defaults
python examples/jax/cutlass_call_export.py
python examples/jax/cutlass_call_export.py --M 512 --N 256
# Run with shape (1024, 512)
python examples/jax/cutlass_call_export.py --M 1024 --N 512
# Export with symbolic shapes.
python examples/jax/cutlass_call_export.py --export_symbolic
"""
import argparse
import cuda.bindings.driver as cuda
import cutlass.cute as cute
import jax
import jax.numpy as jnp
from jax import export
from cutlass.jax import cutlass_call, get_export_disabled_safety_checks, TensorSpec
from cutlass.jax.testing import create_tensor
# Simple element-wise addition kernel: gC[i,j] = gA[i,j] + gB[i,j]
@cute.kernel
def kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
tidx, _, _ = cute.arch.thread_idx()
@@ -84,9 +84,6 @@ def kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
@cute.jit
def launch(stream: cuda.CUstream, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
print("mA: ", mA.layout)
print("mB: ", mB.layout)
print("mC: ", mC.layout)
num_threads_per_block = 256
m, n = mA.shape
kernel(mA, mB, mC).launch(
@@ -96,63 +93,100 @@ def launch(stream: cuda.CUstream, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Ten
)
def run_example(M, N, export_symbolic_shapes):
def _export_and_run(f, ref_f, input_shape_dtype, run_shapes):
"""Export f, serialize/deserialize, then run on each shape in run_shapes.
Both inputs (a, b) are assumed to share the same input_shape_dtype.
"""
print(f"Exporting with input signature: ({input_shape_dtype}, {input_shape_dtype})")
# jax.export can be used to export a jit function containing cutlass_call.
# CUTLASS custom call targets are not on JAX's built-in stable custom-call
# allowlist, so we pass them via disabled_checks to suppress that safety check.
exported = jax.export.export(f, disabled_checks=get_export_disabled_safety_checks())
traced = exported(input_shape_dtype, input_shape_dtype)
blob = traced.serialize()
print(f"Serialized computation is {len(blob)} bytes.")
rehydrated = export.deserialize(blob)
key = jax.random.key(1123)
a_key, b_key = jax.random.split(key, 2)
for shape in run_shapes:
a = create_tensor(shape, dtype=jnp.float32, key=a_key)
b = create_tensor(shape, dtype=jnp.float32, key=b_key)
c = rehydrated.call(a, b)
assert jnp.allclose(c, ref_f(a, b)), f"Mismatch at shape {shape}"
print(f" shape {shape}: OK")
def run_example(M, N):
@jax.jit
def ref_f(a, b):
return jax.nn.sigmoid(a + b)
# The same JAX function is used in all three examples below. The export
# mode is determined entirely by the shape spec passed to jax.export.
@jax.jit
def f(a, b):
call = cutlass_call(launch, output_shape_dtype=a)
return jax.nn.sigmoid(call(a, b))
# ── 1. Concrete shapes ────────────────────────────────────────────────────
# Shapes are fixed constants baked into the export. The deserialized
# computation only accepts exactly these dimensions at runtime.
print("\nConcrete shapes:")
input_shape_dtype = jax.ShapeDtypeStruct((M, N), jnp.float32)
_export_and_run(
f,
ref_f,
input_shape_dtype,
run_shapes=[(M, N)], # concrete exports reject any other shape
)
# ── 2. Unconstrained symbolic shapes ─────────────────────────────────────
# Both dimensions are fully dynamic. The exported computation accepts any
# (M, N) at runtime without recompilation.
print("\nUnconstrained symbolic shapes:")
a_sym, b_sym = export.symbolic_shape("a, b")
input_shape_dtype = jax.ShapeDtypeStruct((a_sym, b_sym), jnp.float32)
_export_and_run(
f,
ref_f,
input_shape_dtype,
run_shapes=[(M, N), (M * 2, N * 4), (M * 4, N * 4)],
)
# ── 3. Constrained symbolic shapes (divisibility) ─────────────────────────
# Shapes are declared as multiples of a tile size via TensorSpec.divisibility.
# The symbolic expression "32*M, 16*N" tells jax.export that dim 0 is always
# a multiple of 32 and dim 1 is always a multiple of 16. This lets the
# compiler generate more efficient code (e.g. no remainder handling).
# Runtime shapes must satisfy these divisibility constraints.
print("\nConstrained symbolic shapes:")
@jax.jit
def ref_f(a, b):
return jax.nn.sigmoid(a + b)
def f_divisible(a, b):
spec = TensorSpec(divisibility=(32, 16))
call = cutlass_call(
launch,
output_shape_dtype=a,
input_spec=(spec, spec),
output_spec=spec,
)
return jax.nn.sigmoid(call(a, b))
# Symbolic or partially shapes are supported by cutlass_call and cute.Tensor
# This allows export of functions calling Cut eDSL kernels w/o having to re-compile
# the kernel for each new shape.
if export_symbolic_shapes:
a, b = export.symbolic_shape("a, b")
export_shape_dtype = jax.ShapeDtypeStruct((a, b), jnp.float32)
else:
export_shape_dtype = jax.ShapeDtypeStruct((M, N), jnp.float32)
print("Exporting with input signature: ")
print(f"({export_shape_dtype}, {export_shape_dtype})")
# jax.export can be used to export a jit function containing cutlass_call.
# The function get_export_disabled_safety_checks() returns a list of custom
# call targets that are used by cutlass_call not part of Jax's built-in
# list of stable custom calls.
exported = jax.export.export(f, disabled_checks=get_export_disabled_safety_checks())
traced = exported(export_shape_dtype, export_shape_dtype)
# Serialize the computation to a byte blob.
blob = traced.serialize()
print(f"Serialized computation is {len(blob)} bytes.")
# Deserialize and run
rehydrated = export.deserialize(blob)
key = jax.random.key(1123)
a_key, b_key = jax.random.split(key, 2)
a = create_tensor((M, N), dtype=jnp.float32, key=a_key)
b = create_tensor((M, N), dtype=jnp.float32, key=b_key)
c = rehydrated.call(a, b)
assert jnp.allclose(c, ref_f(a, b))
# If the computation was exported with dynamic shapes then we can also
# call it with different shapes. The kernel will not be re-compiled
# even though the shapes are changing.
if export_symbolic_shapes:
a = create_tensor((M * 2, N * 4), dtype=jnp.float32, key=a_key)
b = create_tensor((M * 2, N * 4), dtype=jnp.float32, key=b_key)
c = rehydrated.call(a, b)
assert jnp.allclose(c, ref_f(a, b))
a = create_tensor((M * 4, N * 4), dtype=jnp.float32, key=a_key)
b = create_tensor((M * 4, N * 4), dtype=jnp.float32, key=b_key)
c = rehydrated.call(a, b)
assert jnp.allclose(c, ref_f(a, b))
m_sym, n_sym = export.symbolic_shape("32*M, 16*N")
input_shape_dtype = jax.ShapeDtypeStruct((m_sym, n_sym), jnp.float32)
_export_and_run(
f_divisible,
ref_f,
input_shape_dtype,
run_shapes=[(M, N), (M * 2, N * 2), (M * 4, N * 4)],
)
if __name__ == "__main__":
@@ -161,8 +195,7 @@ if __name__ == "__main__":
)
parser.add_argument("--M", default=512, type=int)
parser.add_argument("--N", default=256, type=int)
parser.add_argument("--export_symbolic", action="store_true")
args = parser.parse_args()
run_example(args.M, args.N, args.export_symbolic)
run_example(args.M, args.N)
print("PASS")
examples/python/CuTeDSL/jax/cutlass_call_sharding.py
-2
View File
@@ -27,14 +27,12 @@
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
from functools import partial
import argparse
import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P, AxisType
from jax.experimental.custom_partitioning import custom_partitioning
import cutlass
import cutlass.cute as cute
import cutlass.jax as cjax
from cutlass.jax.testing import create_tensor
examples/python/CuTeDSL/jax/elementwise_apply_example.py
+1 -1
View File
@@ -30,7 +30,7 @@
import argparse
import operator
from functools import partial
from typing import List, Type
from typing import List
import cuda.bindings.driver as cuda
import cutlass
examples/python/CuTeDSL/notebooks/async_pipeline.ipynb
+4 -4
View File
@@ -78,7 +78,7 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -114,7 +114,7 @@
" ]\n",
"\n",
" synced_producer_consumer(SharedStorage, res).launch(\n",
" grid=(1, 1, 1), block=(64, 1, 1), smem=SharedStorage.size_in_bytes()\n",
" grid=(1, 1, 1), block=(64, 1, 1)\n",
" )\n",
"\n",
"\n",
@@ -455,7 +455,7 @@
},
{
"cell_type": "code",
"execution_count": 6,
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
@@ -527,7 +527,7 @@
" ]\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",
" grid=(1, 1, 1), block=(64, 1, 1)\n",
" )\n",
"\n",
"\n",