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

v4.1 release

Browse Source
This commit is contained in:
Junkai-Wu
2025-07-03 20:07:53 +08:00
committed by GitHub
parent b995f93317
commit a1aaf2300a
155 changed files with 18407 additions and 6068 deletions
Download Patch File Download Diff File Expand all files Collapse all files

151
examples/python/CuTeDSL/blackwell/dense_gemm.py
View File

@@ -35,6 +35,7 @@ import torch
import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils.blackwell_helpers as sm100_utils
@@ -211,7 +212,7 @@ class DenseGemmKernel:
self.occupancy = 1
self.threads_per_cta = 128
self.num_smem_capacity = sm100_utils.SMEM_CAPACITY["sm100"]
self.smem_capacity = sm100_utils.SMEM_CAPACITY["sm100"]
def _setup_attributes(self):
"""Set up configurations that are dependent on GEMM inputs
@@ -283,7 +284,7 @@ class DenseGemmKernel:
self.epi_tile,
self.c_dtype,
self.c_layout,
self.num_smem_capacity,
self.smem_capacity,
self.occupancy,
self.use_tma_store,
)
@@ -308,7 +309,7 @@ class DenseGemmKernel:
self.epi_tile,
self.num_c_stage,
)
if cutlass.const_expr(self.use_tma_store)
if self.use_tma_store
else None
)
@@ -372,9 +373,11 @@ class DenseGemmKernel:
atom_thr_size = cute.size(tiled_mma.thr_id.shape)
# Setup TMA load for A
a_op = self._get_tma_atom_kind(atom_thr_size, self.is_a_mcast)
a_op = sm100_utils.cluster_shape_to_tma_atom_A(
self.cluster_shape_mn, tiled_mma.thr_id
)
a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
tma_atom_a, tma_tensor_a = cute.nvgpu.make_tma_tile_atom_A(
tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
a_op,
a,
a_smem_layout,
@@ -387,9 +390,11 @@ class DenseGemmKernel:
)
# Setup TMA load for B
b_op = self._get_tma_atom_kind(atom_thr_size, self.is_b_mcast)
b_op = sm100_utils.cluster_shape_to_tma_atom_B(
self.cluster_shape_mn, tiled_mma.thr_id
)
b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
tma_atom_b, tma_tensor_b = cute.nvgpu.make_tma_tile_atom_B(
tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
b_op,
b,
b_smem_layout,
@@ -413,7 +418,7 @@ class DenseGemmKernel:
cute.make_identity_layout(c.shape), self.epi_tile
)
epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0))
tma_atom_c, tma_tensor_c = cpasync.make_tma_tile_atom(
tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
cpasync.CopyBulkTensorTileS2GOp(),
c,
epi_smem_layout,
@@ -426,9 +431,7 @@ class DenseGemmKernel:
self.buffer_align_bytes = 1024
c_smem_size = (
cute.cosize(self.c_smem_layout_staged.outer)
if cutlass.const_expr(self.use_tma_store)
else 0
cute.cosize(self.c_smem_layout_staged.outer) if self.use_tma_store else 0
)
# Define shared storage for kernel
@@ -472,7 +475,7 @@ class DenseGemmKernel:
tma_atom_b,
tma_tensor_b,
tma_atom_c,
tma_tensor_c if cutlass.const_expr(self.use_tma_store) else c,
tma_tensor_c if self.use_tma_store else c,
self.cluster_layout_vmnk,
self.a_smem_layout_staged,
self.b_smem_layout_staged,
@@ -556,12 +559,12 @@ class DenseGemmKernel:
tmem_holding_buf = storage.tmem_holding_buf
# Initialize mainloop ab_pipeline (barrier) and states
ab_pipeline_producer_group = utils.CooperativeGroup(utils.Agent.Thread)
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
ab_pipeline_consumer_group = utils.CooperativeGroup(
utils.Agent.Thread, num_tma_producer
ab_pipeline_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, num_tma_producer
)
ab_pipeline = utils.PipelineTmaUmma.create(
ab_pipeline = pipeline.PipelineTmaUmma.create(
barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
num_stages=self.num_ab_stage,
producer_group=ab_pipeline_producer_group,
@@ -569,30 +572,30 @@ class DenseGemmKernel:
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cluster_layout_vmnk,
)
ab_producer_state = utils.make_pipeline_state(
utils.PipelineUserType.Producer, self.num_ab_stage
ab_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_ab_stage
)
ab_consumer_state = utils.make_pipeline_state(
utils.PipelineUserType.Consumer, self.num_ab_stage
ab_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_ab_stage
)
# Initialize acc_pipeline (barrier) and states
acc_pipeline_producer_group = utils.CooperativeGroup(utils.Agent.Thread)
acc_pipeline_consumer_group = utils.CooperativeGroup(
utils.Agent.Thread, self.threads_per_cta, self.threads_per_cta
acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
acc_pipeline_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, self.threads_per_cta, self.threads_per_cta
)
acc_pipeline = utils.PipelineUmmaAsync.create(
acc_pipeline = pipeline.PipelineUmmaAsync.create(
barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
num_stages=self.num_acc_stage,
producer_group=acc_pipeline_producer_group,
consumer_group=acc_pipeline_consumer_group,
cta_layout_vmnk=cluster_layout_vmnk,
)
acc_producer_state = utils.make_pipeline_state(
utils.PipelineUserType.Producer, self.num_acc_stage
acc_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_acc_stage
)
acc_consumer_state = utils.make_pipeline_state(
utils.PipelineUserType.Consumer, self.num_acc_stage
acc_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
)
# Tensor memory dealloc barrier init
@@ -600,7 +603,7 @@ class DenseGemmKernel:
if warp_idx == 0:
num_tmem_dealloc_threads = 32
with cute.arch.elect_one():
cute.arch.mbarrier_init_arrive_cnt(
cute.arch.mbarrier_init(
tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads
)
cute.arch.mbarrier_init_fence()
@@ -617,7 +620,7 @@ class DenseGemmKernel:
storage.sC.get_tensor(
c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner
)
if cutlass.const_expr(self.use_tma_store)
if self.use_tma_store
else None
)
# (MMA, MMA_M, MMA_K, STAGE)
@@ -634,7 +637,7 @@ class DenseGemmKernel:
#
a_full_mcast_mask = None
b_full_mcast_mask = None
if self.is_a_mcast or self.is_b_mcast or use_2cta_instrs:
if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs):
a_full_mcast_mask = cpasync.create_tma_multicast_mask(
cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
)
@@ -645,15 +648,15 @@ class DenseGemmKernel:
#
# Local_tile partition global tensors
#
# (bM, bK, loopM, loopK, loopL)
# (bM, bK, RestM, RestK, RestL)
gA_mkl = cute.local_tile(
mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
# (bN, bK, loopN, loopK, loopL)
# (bN, bK, RestN, RestK, RestL)
gB_nkl = cute.local_tile(
mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
# (bM, bN, loopM, loopN, loopL)
# (bM, bN, RestM, RestN, RestL)
gC_mnl = cute.local_tile(
mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
)
@@ -663,11 +666,11 @@ class DenseGemmKernel:
# Partition global tensor for TiledMMA_A/B/C
#
thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
# (MMA, MMA_M, MMA_K, loopM, loopK, loopL)
# (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
tCgA = thr_mma.partition_A(gA_mkl)
# (MMA, MMA_N, MMA_K, loopN, loopK, loopL)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgB = thr_mma.partition_B(gB_nkl)
# (MMA, MMA_M, MMA_N, loopM, loopN, loopL)
# (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
tCgC = thr_mma.partition_C(gC_mnl)
#
@@ -678,7 +681,7 @@ class DenseGemmKernel:
cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
)
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), loopM, loopK, loopL)
# ((atom_v, rest_v), RestM, RestK, RestL)
tAsA, tAgA = cpasync.tma_partition(
tma_atom_a,
block_in_cluster_coord_vmnk[2],
@@ -691,7 +694,7 @@ class DenseGemmKernel:
cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
)
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), loopN, loopK, loopL)
# ((atom_v, rest_v), RestN, RestK, RestL)
tBsB, tBgB = cpasync.tma_partition(
tma_atom_b,
block_in_cluster_coord_vmnk[1],
@@ -771,9 +774,9 @@ class DenseGemmKernel:
#
# Slice to per mma tile index
#
# ((atom_v, rest_v), loopK)
# ((atom_v, rest_v), RestK)
tAgA = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
# ((atom_v, rest_v), loopK)
# ((atom_v, rest_v), RestK)
tBgB = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
if cutlass.const_expr(self.use_tma_store):
# ((ATOM_V, REST_V), EPI_M, EPI_N)
@@ -797,7 +800,7 @@ class DenseGemmKernel:
#
# Prefetch TMA load A/B
#
for prefetch_idx in cutlass.range_dynamic(prefetch_k_block_cnt, unroll=1):
for prefetch_idx in cutlass.range(prefetch_k_block_cnt, unroll=1):
# Conditionally wait for AB buffer empty
ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status)
@@ -833,7 +836,7 @@ class DenseGemmKernel:
#
# MMA mainloop
#
for k_block in cutlass.range_dynamic(0, k_block_cnt, 1, unroll=1):
for k_block in range(k_block_cnt):
# Conditionally wait for AB buffer empty
ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status)
@@ -860,7 +863,7 @@ class DenseGemmKernel:
# tCtAcc += tCrA * tCrB
num_kphases = cute.size(tCrA, mode=[2])
for kphase_idx in range(num_kphases):
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
kphase_coord = (None, None, kphase_idx, ab_consumer_state.index)
cute.gemm(
@@ -917,10 +920,10 @@ class DenseGemmKernel:
c_pipeline = None
if cutlass.const_expr(self.use_tma_store):
# Initialize tma store c_pipeline
c_producer_group = utils.CooperativeGroup(
utils.Agent.Thread, self.threads_per_cta, self.threads_per_cta
c_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, self.threads_per_cta, self.threads_per_cta
)
c_pipeline = utils.PipelineTmaStore.create(
c_pipeline = pipeline.PipelineTmaStore.create(
num_stages=self.num_c_stage,
producer_group=c_producer_group,
)
@@ -929,7 +932,7 @@ class DenseGemmKernel:
# Store accumulator to global memory in subtiles
#
subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
for subtile_idx in cutlass.range_dynamic(subtile_cnt):
for subtile_idx in range(subtile_cnt):
#
# Load accumulator from tensor memory buffer to register
#
@@ -1007,7 +1010,7 @@ class DenseGemmKernel:
#
if warp_idx == 0:
# Reverse prefetch_k_block_cnt times to next available buffer
for i in cutlass.range_dynamic(prefetch_k_block_cnt):
for i in range(prefetch_k_block_cnt):
ab_producer_state.reverse()
ab_pipeline.producer_tail(ab_producer_state)
return
@@ -1063,11 +1066,11 @@ class DenseGemmKernel:
# (T2R, T2R_M, T2R_N, EPI_M, EPI_M)
tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL)
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL)
gC_mnl_epi = cute.flat_divide(
gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
)
# (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL)
# (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(
@@ -1149,7 +1152,7 @@ class DenseGemmKernel:
- tTR_gC: The partitioned global tensor C
:rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]
"""
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL)
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL)
gC_epi = cute.flat_divide(
gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
)
@@ -1158,7 +1161,7 @@ class DenseGemmKernel:
sC_for_tma_partition = cute.group_modes(sC, 0, 2)
gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2)
# ((ATOM_V, REST_V), EPI_M, EPI_N)
# ((ATOM_V, REST_V), EPI_M, EPI_N, loopM, loopN, loopL)
# ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL)
bSG_sC, bSG_gC = cpasync.tma_partition(
tma_atom_c,
0,
@@ -1169,7 +1172,7 @@ class DenseGemmKernel:
return tma_atom_c, bSG_sC, bSG_gC
else:
tiled_copy_t2r = atom
# (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL)
# (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL)
thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
tTR_gC = thr_copy_t2r.partition_D(gC_epi)
# (T2R, T2R_M, T2R_N)
@@ -1188,7 +1191,7 @@ class DenseGemmKernel:
epi_tile: cute.Tile,
c_dtype: Type[cutlass.Numeric],
c_layout: utils.LayoutEnum,
num_smem_capacity: int,
smem_capacity: int,
occupancy: int,
use_tma_store: bool,
) -> Tuple[int, int, int]:
@@ -1208,8 +1211,8 @@ class DenseGemmKernel:
:type c_dtype: type[cutlass.Numeric]
:param c_layout: Layout enum of operand C in global memory.
:type c_layout: utils.LayoutEnum
:param num_smem_capacity: Total available shared memory capacity in bytes.
:type num_smem_capacity: int
:param smem_capacity: Total available shared memory capacity in bytes.
:type smem_capacity: int
:param occupancy: Target number of CTAs per SM (occupancy).
:type occupancy: int
:param use_tma_store: Whether TMA store is enabled.
@@ -1263,7 +1266,7 @@ class DenseGemmKernel:
# Subtract reserved bytes and initial C stages bytes
# Divide remaining by bytes needed per A/B stage
num_ab_stage = (
num_smem_capacity - (occupancy + 1) * (mbar_helpers_bytes + c_bytes)
smem_capacity - (occupancy + 1) * (mbar_helpers_bytes + c_bytes)
) // ab_bytes_per_stage
# Refine epilogue stages:
@@ -1271,7 +1274,7 @@ class DenseGemmKernel:
# Add remaining unused smem to epilogue
if use_tma_store:
num_c_stage += (
num_smem_capacity
smem_capacity
- ab_bytes_per_stage * num_ab_stage
- (occupancy + 1) * (mbar_helpers_bytes + c_bytes)
) // ((occupancy + 1) * c_bytes_per_stage)
@@ -1309,36 +1312,6 @@ class DenseGemmKernel:
return grid
@staticmethod
def _get_tma_atom_kind(
atom_sm_cnt: cutlass.Int32, mcast: cutlass.Boolean
) -> Union[
cpasync.CopyBulkTensorTileG2SMulticastOp, cpasync.CopyBulkTensorTileG2SOp
]:
"""
Select the appropriate TMA copy atom based on the number of SMs and the multicast flag.
:param atom_sm_cnt: The number of SMs
:type atom_sm_cnt: cutlass.Int32
:param mcast: The multicast flag
:type mcast: cutlass.Boolean
:return: The appropriate TMA copy atom kind
:rtype: cpasync.CopyBulkTensorTileG2SMulticastOp or cpasync.CopyBulkTensorTileG2SOp
:raise ValueError: If the atom_sm_cnt is invalid
"""
if atom_sm_cnt == 2 and mcast:
return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.TWO)
elif atom_sm_cnt == 2 and not mcast:
return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.TWO)
elif atom_sm_cnt == 1 and mcast:
return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.ONE)
elif atom_sm_cnt == 1 and not mcast:
return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE)
raise ValueError(f"Invalid atom_sm_cnt: {atom_sm_cnt} and {mcast}")
@staticmethod
def _compute_num_tmem_alloc_cols(
tiled_mma: cute.TiledMma, mma_tiler: Tuple[int, int, int]

172
examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
View File

@@ -37,6 +37,7 @@ import cutlass.cute as cute
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.pipeline as pipeline
import cutlass.utils.blackwell_helpers as sm100_utils
from cutlass.cute.runtime import from_dlpack
@@ -225,7 +226,7 @@ class PersistentDenseGemmKernel:
self.cta_sync_bar_id = 0
self.epilog_sync_bar_id = 1
self.tmem_ptr_sync_bar_id = 2
self.num_smem_capacity = sm100_utils.SMEM_CAPACITY["sm100"]
self.smem_capacity = sm100_utils.SMEM_CAPACITY["sm100"]
def _setup_attributes(self):
"""Set up configurations that are dependent on GEMM inputs
@@ -297,7 +298,7 @@ class PersistentDenseGemmKernel:
self.epi_tile,
self.c_dtype,
self.c_layout,
self.num_smem_capacity,
self.smem_capacity,
self.occupancy,
self.use_tma_store,
)
@@ -389,9 +390,11 @@ class PersistentDenseGemmKernel:
atom_thr_size = cute.size(tiled_mma.thr_id.shape)
# Setup TMA load for A
a_op = self._get_tma_atom_kind(atom_thr_size, self.is_a_mcast)
a_op = sm100_utils.cluster_shape_to_tma_atom_A(
self.cluster_shape_mn, tiled_mma.thr_id
)
a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
tma_atom_a, tma_tensor_a = cute.nvgpu.make_tma_tile_atom_A(
tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
a_op,
a,
a_smem_layout,
@@ -404,9 +407,11 @@ class PersistentDenseGemmKernel:
)
# Setup TMA load for B
b_op = self._get_tma_atom_kind(atom_thr_size, self.is_b_mcast)
b_op = sm100_utils.cluster_shape_to_tma_atom_B(
self.cluster_shape_mn, tiled_mma.thr_id
)
b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
tma_atom_b, tma_tensor_b = cute.nvgpu.make_tma_tile_atom_B(
tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
b_op,
b,
b_smem_layout,
@@ -430,7 +435,7 @@ class PersistentDenseGemmKernel:
cute.make_identity_layout(c.shape), self.epi_tile
)
epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0))
tma_atom_c, tma_tensor_c = cpasync.make_tma_tile_atom(
tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
cpasync.CopyBulkTensorTileS2GOp(),
c,
epi_smem_layout,
@@ -571,12 +576,12 @@ class PersistentDenseGemmKernel:
tmem_holding_buf = storage.tmem_holding_buf
# Initialize mainloop ab_pipeline (barrier) and states
ab_pipeline_producer_group = utils.CooperativeGroup(utils.Agent.Thread)
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
ab_pipeline_consumer_group = utils.CooperativeGroup(
utils.Agent.Thread, num_tma_producer
ab_pipeline_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, num_tma_producer
)
ab_pipeline = utils.PipelineTmaUmma.create(
ab_pipeline = pipeline.PipelineTmaUmma.create(
barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
num_stages=self.num_ab_stage,
producer_group=ab_pipeline_producer_group,
@@ -586,14 +591,14 @@ class PersistentDenseGemmKernel:
)
# Initialize acc_pipeline (barrier) and states
acc_pipeline_producer_group = utils.CooperativeGroup(utils.Agent.Thread)
acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
num_acc_consumer_threads = len(self.epilog_warp_id) * (
2 if use_2cta_instrs else 1
)
acc_pipeline_consumer_group = utils.CooperativeGroup(
utils.Agent.Thread, num_acc_consumer_threads
acc_pipeline_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, num_acc_consumer_threads
)
acc_pipeline = utils.PipelineUmmaAsync.create(
acc_pipeline = pipeline.PipelineUmmaAsync.create(
barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
num_stages=self.num_acc_stage,
producer_group=acc_pipeline_producer_group,
@@ -606,7 +611,7 @@ class PersistentDenseGemmKernel:
if warp_idx == self.tma_warp_id:
num_tmem_dealloc_threads = 32
with cute.arch.elect_one():
cute.arch.mbarrier_init_arrive_cnt(
cute.arch.mbarrier_init(
tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads
)
cute.arch.mbarrier_init_fence()
@@ -640,7 +645,7 @@ class PersistentDenseGemmKernel:
#
a_full_mcast_mask = None
b_full_mcast_mask = None
if self.is_a_mcast or self.is_b_mcast or use_2cta_instrs:
if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs):
a_full_mcast_mask = cpasync.create_tma_multicast_mask(
cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
)
@@ -651,15 +656,15 @@ class PersistentDenseGemmKernel:
#
# Local_tile partition global tensors
#
# (bM, bK, loopM, loopK, loopL)
# (bM, bK, RestM, RestK, RestL)
gA_mkl = cute.local_tile(
mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
# (bN, bK, loopN, loopK, loopL)
# (bN, bK, RestN, RestK, RestL)
gB_nkl = cute.local_tile(
mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
# (bM, bN, loopM, loopN, loopL)
# (bM, bN, RestM, RestN, RestL)
gC_mnl = cute.local_tile(
mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
)
@@ -669,11 +674,11 @@ class PersistentDenseGemmKernel:
# Partition global tensor for TiledMMA_A/B/C
#
thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
# (MMA, MMA_M, MMA_K, loopM, loopK, loopL)
# (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
tCgA = thr_mma.partition_A(gA_mkl)
# (MMA, MMA_N, MMA_K, loopN, loopK, loopL)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgB = thr_mma.partition_B(gB_nkl)
# (MMA, MMA_M, MMA_N, loopM, loopN, loopL)
# (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
tCgC = thr_mma.partition_C(gC_mnl)
#
@@ -684,7 +689,7 @@ class PersistentDenseGemmKernel:
cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
)
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), loopM, loopK, loopL)
# ((atom_v, rest_v), RestM, RestK, RestL)
tAsA, tAgA = cpasync.tma_partition(
tma_atom_a,
block_in_cluster_coord_vmnk[2],
@@ -697,7 +702,7 @@ class PersistentDenseGemmKernel:
cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
)
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), loopM, loopK, loopL)
# ((atom_v, rest_v), RestM, RestK, RestL)
tBsB, tBgB = cpasync.tma_partition(
tma_atom_b,
block_in_cluster_coord_vmnk[1],
@@ -743,12 +748,11 @@ class PersistentDenseGemmKernel:
)
work_tile = tile_sched.initial_work_tile_info()
ab_producer_state = utils.make_pipeline_state(
utils.PipelineUserType.Producer, self.num_ab_stage
ab_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_ab_stage
)
while work_tile.is_valid_tile:
# Get tile coord from tile scheduler
cur_tile_coord = work_tile.tile_idx
mma_tile_coord_mnl = (
@@ -760,11 +764,11 @@ class PersistentDenseGemmKernel:
#
# Slice to per mma tile index
#
# ((atom_v, rest_v), loopK)
# ((atom_v, rest_v), RestK)
tAgA_slice = tAgA[
(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])
]
# ((atom_v, rest_v), loopK)
# ((atom_v, rest_v), RestK)
tBgB_slice = tBgB[
(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])
]
@@ -779,7 +783,7 @@ class PersistentDenseGemmKernel:
#
# Tma load loop
#
for k_block in cutlass.range_dynamic(0, k_block_cnt, 1, unroll=1):
for k_block in cutlass.range(0, k_block_cnt, 1, unroll=1):
# Conditionally wait for AB buffer empty
ab_pipeline.producer_acquire(
ab_producer_state, peek_ab_empty_status
@@ -852,15 +856,14 @@ class PersistentDenseGemmKernel:
)
work_tile = tile_sched.initial_work_tile_info()
ab_consumer_state = utils.make_pipeline_state(
utils.PipelineUserType.Consumer, self.num_ab_stage
ab_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_ab_stage
)
acc_producer_state = utils.make_pipeline_state(
utils.PipelineUserType.Producer, self.num_acc_stage
acc_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_acc_stage
)
while work_tile.is_valid_tile:
# Get tile coord from tile scheduler
cur_tile_coord = work_tile.tile_idx
mma_tile_coord_mnl = (
@@ -895,7 +898,7 @@ class PersistentDenseGemmKernel:
#
# Mma mainloop
#
for k_block in cutlass.range_dynamic(0, k_block_cnt, 1, unroll=1):
for k_block in range(k_block_cnt):
if is_leader_cta:
# Conditionally wait for AB buffer full
ab_pipeline.consumer_wait(
@@ -904,7 +907,7 @@ class PersistentDenseGemmKernel:
# tCtAcc += tCrA * tCrB
num_kphases = cute.size(tCrA, mode=[2])
for kphase_idx in range(num_kphases):
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
kphase_coord = (
None,
None,
@@ -989,10 +992,12 @@ class PersistentDenseGemmKernel:
# 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 = None
@@ -1008,16 +1013,20 @@ class PersistentDenseGemmKernel:
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
)
else:
simt_atom, tTR_rC, tTR_gC_partitioned = (
self.epilog_gmem_copy_and_partition(
epi_tidx, tiled_copy_t2r, tCgC, epi_tile, sC
)
(
simt_atom,
tTR_rC,
tTR_gC_partitioned,
) = self.epilog_gmem_copy_and_partition(
epi_tidx, tiled_copy_t2r, tCgC, epi_tile, sC
)
#
@@ -1028,25 +1037,24 @@ class PersistentDenseGemmKernel:
)
work_tile = tile_sched.initial_work_tile_info()
acc_consumer_state = utils.make_pipeline_state(
utils.PipelineUserType.Consumer, self.num_acc_stage
acc_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
)
c_pipeline = None
if cutlass.const_expr(self.use_tma_store):
# Threads/warps participating in tma store pipeline
c_producer_group = utils.CooperativeGroup(
utils.Agent.Thread,
c_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread,
32 * len(self.epilog_warp_id),
32 * len(self.epilog_warp_id),
)
c_pipeline = utils.PipelineTmaStore.create(
c_pipeline = pipeline.PipelineTmaStore.create(
num_stages=self.num_c_stage,
producer_group=c_producer_group,
)
while work_tile.is_valid_tile:
# Get tile coord from tile scheduler
cur_tile_coord = work_tile.tile_idx
mma_tile_coord_mnl = (
@@ -1105,7 +1113,7 @@ class PersistentDenseGemmKernel:
#
subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt
for subtile_idx in cutlass.range_dynamic(subtile_cnt):
for subtile_idx in cutlass.range(subtile_cnt):
#
# Load accumulator from tensor memory buffer to register
#
@@ -1259,11 +1267,11 @@ class PersistentDenseGemmKernel:
# (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE)
tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL)
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL)
gC_mnl_epi = cute.flat_divide(
gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
)
# (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL)
# (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(
@@ -1346,7 +1354,7 @@ class PersistentDenseGemmKernel:
- tTR_gC: The partitioned global tensor C
:rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]
"""
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL)
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL)
gC_epi = cute.flat_divide(
gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
)
@@ -1355,7 +1363,7 @@ class PersistentDenseGemmKernel:
sC_for_tma_partition = cute.group_modes(sC, 0, 2)
gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2)
# ((ATOM_V, REST_V), EPI_M, EPI_N)
# ((ATOM_V, REST_V), EPI_M, EPI_N, loopM, loopN, loopL)
# ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL)
bSG_sC, bSG_gC = cpasync.tma_partition(
tma_atom_c,
0,
@@ -1366,7 +1374,7 @@ class PersistentDenseGemmKernel:
return tma_atom_c, bSG_sC, bSG_gC
else:
tiled_copy_t2r = atom
# (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL)
# (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL)
thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
tTR_gC = thr_copy_t2r.partition_D(gC_epi)
# (T2R, T2R_M, T2R_N)
@@ -1385,7 +1393,7 @@ class PersistentDenseGemmKernel:
epi_tile: cute.Tile,
c_dtype: Type[cutlass.Numeric],
c_layout: utils.LayoutEnum,
num_smem_capacity: int,
smem_capacity: int,
occupancy: int,
use_tma_store: bool,
) -> Tuple[int, int, int]:
@@ -1405,8 +1413,8 @@ class PersistentDenseGemmKernel:
:type c_dtype: type[cutlass.Numeric]
:param c_layout: Layout enum of operand C.
:type c_layout: utils.LayoutEnum
:param num_smem_capacity: Total available shared memory capacity in bytes.
:type num_smem_capacity: int
:param smem_capacity: Total available shared memory capacity in bytes.
:type smem_capacity: int
:param occupancy: Target number of CTAs per SM (occupancy).
:type occupancy: int
:param use_tma_store: Whether TMA store is enabled.
@@ -1461,7 +1469,7 @@ class PersistentDenseGemmKernel:
# Subtract reserved bytes and initial C stages bytes
# Divide remaining by bytes needed per A/B stage
num_ab_stage = (
num_smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes)
smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes)
) // ab_bytes_per_stage
# Refine epilogue stages:
@@ -1469,7 +1477,7 @@ class PersistentDenseGemmKernel:
# Add remaining unused smem to epilogue
if use_tma_store:
num_c_stage += (
num_smem_capacity
smem_capacity
- occupancy * ab_bytes_per_stage * num_ab_stage
- occupancy * (mbar_helpers_bytes + c_bytes)
) // (occupancy * c_bytes_per_stage)
@@ -1512,36 +1520,6 @@ class PersistentDenseGemmKernel:
return tile_sched_params, grid
@staticmethod
def _get_tma_atom_kind(
atom_sm_cnt: cutlass.Int32, mcast: cutlass.Boolean
) -> Union[
cpasync.CopyBulkTensorTileG2SMulticastOp, cpasync.CopyBulkTensorTileG2SOp
]:
"""
Select the appropriate TMA copy atom based on the number of SMs and the multicast flag.
:param atom_sm_cnt: The number of SMs
:type atom_sm_cnt: cutlass.Int32
:param mcast: The multicast flag
:type mcast: cutlass.Boolean
:return: The appropriate TMA copy atom kind
:rtype: cpasync.CopyBulkTensorTileG2SMulticastOp or cpasync.CopyBulkTensorTileG2SOp
:raise ValueError: If the atom_sm_cnt is invalid
"""
if atom_sm_cnt == 2 and mcast:
return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.TWO)
elif atom_sm_cnt == 2 and not mcast:
return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.TWO)
elif atom_sm_cnt == 1 and mcast:
return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.ONE)
elif atom_sm_cnt == 1 and not mcast:
return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE)
raise ValueError(f"Invalid atom_sm_cnt: {atom_sm_cnt} and {mcast}")
@staticmethod
def _compute_num_tmem_alloc_cols(
tiled_mma: cute.TiledMma,

1852
examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py Normal file
View File

File diff suppressed because it is too large Load Diff

2235
examples/python/CuTeDSL/blackwell/fmha.py
View File

File diff suppressed because it is too large Load Diff

340
examples/python/CuTeDSL/blackwell/grouped_gemm.py
View File

@@ -40,7 +40,6 @@ import cutlass.utils as utils
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.torch as cutlass_torch
from cutlass.cute.runtime import from_dlpack
"""
A grouped GEMM example for the NVIDIA Blackwell SM100 architecture using CUTE DSL
@@ -89,7 +88,6 @@ there are also the following constrains:
class GroupedGemmKernel:
def __init__(
self,
acc_dtype: type[cutlass.Numeric],
@@ -159,7 +157,7 @@ class GroupedGemmKernel:
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
self.num_smem_capacity = sm100_utils.SMEM_CAPACITY["sm100"]
self.smem_capacity = sm100_utils.SMEM_CAPACITY["sm100"]
self.num_tma_load_bytes = 0
def _setup_attributes(self):
@@ -217,18 +215,20 @@ class GroupedGemmKernel:
)
# Setup A/B/C stage count in shared memory and ACC stage count in tensor memory
self.num_acc_stage, self.num_ab_stage, self.num_epi_stage = (
self._compute_stages(
tiled_mma,
self.mma_tiler,
self.a_dtype,
self.b_dtype,
self.epi_tile,
self.c_dtype,
self.c_layout,
self.num_smem_capacity,
self.occupancy,
)
(
self.num_acc_stage,
self.num_ab_stage,
self.num_epi_stage,
) = self._compute_stages(
tiled_mma,
self.mma_tiler,
self.a_dtype,
self.b_dtype,
self.epi_tile,
self.c_dtype,
self.c_layout,
self.smem_capacity,
self.occupancy,
)
self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
@@ -355,9 +355,11 @@ class GroupedGemmKernel:
atom_thr_size = cute.size(tiled_mma.thr_id.shape)
# Setup TMA load for A
a_op = self._get_tma_atom_kind(atom_thr_size, self.is_a_mcast)
a_op = sm100_utils.cluster_shape_to_tma_atom_A(
self.cluster_shape_mn, tiled_mma.thr_id
)
a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
tma_atom_a, tma_tensor_a = cute.nvgpu.make_tma_tile_atom_A(
tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
a_op,
initial_a,
a_smem_layout,
@@ -367,9 +369,11 @@ class GroupedGemmKernel:
)
# Setup TMA load for B
b_op = self._get_tma_atom_kind(atom_thr_size, self.is_b_mcast)
b_op = sm100_utils.cluster_shape_to_tma_atom_B(
self.cluster_shape_mn, tiled_mma.thr_id
)
b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
tma_atom_b, tma_tensor_b = cute.nvgpu.make_tma_tile_atom_B(
tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
b_op,
initial_b,
b_smem_layout,
@@ -389,7 +393,7 @@ class GroupedGemmKernel:
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_tma_tile_atom(
tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
cpasync.CopyBulkTensorTileS2GOp(),
initial_c,
epi_smem_layout,
@@ -403,9 +407,7 @@ class GroupedGemmKernel:
self.buffer_align_bytes = 1024
self.size_tensormap_in_i64 = (
0
if cutlass.const_expr(
self.tensormap_update_mode == utils.TensorMapUpdateMode.GMEM
)
if self.tensormap_update_mode == utils.TensorMapUpdateMode.GMEM
else GroupedGemmKernel.num_tensormaps
* GroupedGemmKernel.bytes_per_tensormap
// 8
@@ -564,16 +566,16 @@ class GroupedGemmKernel:
for k_stage in range(self.num_ab_stage):
num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
with cute.arch.elect_one():
cute.arch.mbarrier_init_arrive_cnt(ab_full_mbar_ptr + k_stage, 1)
cute.arch.mbarrier_init_arrive_cnt(
cute.arch.mbarrier_init(ab_full_mbar_ptr + k_stage, 1)
cute.arch.mbarrier_init(
ab_empty_mbar_ptr + k_stage, num_tma_producer
)
# Accumulator barrier init
if warp_idx == self.mma_warp_id:
for acc_stage in range(self.num_acc_stage):
with cute.arch.elect_one():
cute.arch.mbarrier_init_arrive_cnt(acc_full_mbar_ptr + acc_stage, 1)
cute.arch.mbarrier_init_arrive_cnt(
cute.arch.mbarrier_init(acc_full_mbar_ptr + acc_stage, 1)
cute.arch.mbarrier_init(
acc_empty_mbar_ptr + acc_stage, 8 if use_2cta_instrs else 4
)
# Tensor memory dealloc barrier init
@@ -581,7 +583,7 @@ class GroupedGemmKernel:
if warp_idx == self.tma_warp_id:
num_tmem_dealloc_threads = 32
with cute.arch.elect_one():
cute.arch.mbarrier_init_arrive_cnt(
cute.arch.mbarrier_init(
tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads
)
cute.arch.mbarrier_init_fence()
@@ -612,7 +614,7 @@ class GroupedGemmKernel:
a_full_mcast_mask = None
b_full_mcast_mask = None
ab_empty_mcast_mask = None
if self.is_a_mcast or self.is_b_mcast or use_2cta_instrs:
if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs):
a_full_mcast_mask = cpasync.create_tma_multicast_mask(
cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
)
@@ -621,7 +623,7 @@ class GroupedGemmKernel:
)
ab_empty_mcast_mask = a_full_mcast_mask | b_full_mcast_mask
acc_full_mcast_mask = None
if use_2cta_instrs:
if cutlass.const_expr(use_2cta_instrs):
acc_full_mcast_mask = cute.make_layout_image_mask(
cluster_layout_vmnk, block_in_cluster_coord_vmnk, mode=0
)
@@ -646,15 +648,15 @@ class GroupedGemmKernel:
#
# Local_tile partition global tensors
#
# (bM, bK, loopM, loopK, loopL)
# (bM, bK, RestM, RestK, RestL)
gA_mkl = cute.local_tile(
mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
# (bN, bK, loopN, loopK, loopL)
# (bN, bK, RestN, RestK, RestL)
gB_nkl = cute.local_tile(
mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
# (bM, bN, loopM, loopN, loopL)
# (bM, bN, RestM, RestN, RestL)
gC_mnl = cute.local_tile(
mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
)
@@ -663,11 +665,11 @@ class GroupedGemmKernel:
# Partition global tensor for TiledMMA_A/B/C
#
thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
# (MMA, MMA_M, MMA_K, loopM, loopK, loopL)
# (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
tCgA = thr_mma.partition_A(gA_mkl)
# (MMA, MMA_N, MMA_K, loopN, loopK, loopL)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgB = thr_mma.partition_B(gB_nkl)
# (MMA, MMA_M, MMA_N, loopM, loopN, loopL)
# (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
tCgC = thr_mma.partition_C(gC_mnl)
#
@@ -677,7 +679,7 @@ class GroupedGemmKernel:
cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
)
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), loopM, loopK, loopL)
# ((atom_v, rest_v), RestM, RestK, RestL)
tAsA, tAgA = cpasync.tma_partition(
tma_atom_a,
block_in_cluster_coord_vmnk[2],
@@ -690,7 +692,7 @@ class GroupedGemmKernel:
cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
)
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), loopM, loopK, loopL)
# ((atom_v, rest_v), RestM, RestK, RestL)
tBsB, tBgB = cpasync.tma_partition(
tma_atom_b,
block_in_cluster_coord_vmnk[1],
@@ -849,11 +851,11 @@ class GroupedGemmKernel:
#
# Slice to per mma tile index
#
# ((atom_v, rest_v), loopK)
# ((atom_v, rest_v), RestK)
tAgA_slice = tAgA[
(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])
]
# ((atom_v, rest_v), loopK)
# ((atom_v, rest_v), RestK)
tBgB_slice = tBgB[
(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])
]
@@ -867,7 +869,7 @@ class GroupedGemmKernel:
tma_wr_ab_empty_phase = (
num_prev_k_blk + tma_wr_k_block
) // self.num_ab_stage % 2 ^ 1
peek_ab_empty_status = cute.arch.conditional_mbarrier_try_wait(
peek_ab_empty_status = cute.arch.mbarrier_conditional_try_wait(
tma_wr_k_block < cur_k_block_cnt,
ab_empty_mbar_ptr + smem_wr_buffer,
tma_wr_ab_empty_phase,
@@ -879,7 +881,7 @@ class GroupedGemmKernel:
#
# Tma load loop
#
for k_block in cutlass.range_dynamic(0, cur_k_block_cnt, 1, unroll=1):
for k_block in cutlass.range(0, cur_k_block_cnt, 1, unroll=1):
tma_wr_k_block_next = tma_wr_k_block + 1
smem_wr_buffer_next = (
num_prev_k_blk + tma_wr_k_block_next
@@ -898,10 +900,10 @@ class GroupedGemmKernel:
ab_empty_mbar_ptr + smem_wr_buffer, tma_wr_ab_empty_phase
)
# Init AB buffer full transaction byte
# Arrive AB buffer and expect full transaction bytes
if is_leader_cta:
with cute.arch.elect_one():
cute.arch.mbarrier_init_tx_bytes(
cute.arch.mbarrier_arrive_and_expect_tx(
smem_full_mbar_ptr, self.num_tma_load_bytes
)
@@ -930,7 +932,7 @@ class GroupedGemmKernel:
)
# Peek (try_wait) AB buffer empty for k_block = prefetch_k_block_cnt + k_block + 1
peek_ab_empty_status = cute.arch.conditional_mbarrier_try_wait(
peek_ab_empty_status = cute.arch.mbarrier_conditional_try_wait(
tma_wr_k_block_next < cur_k_block_cnt,
ab_empty_mbar_ptr + smem_wr_buffer_next,
tma_wr_ab_empty_phase_next,
@@ -999,11 +1001,12 @@ class GroupedGemmKernel:
while work_tile.is_valid_tile:
cur_tile_coord = work_tile.tile_idx
# MMA warp is only interested in number of tiles along K dimension
cur_k_block_cnt, cur_group_idx = (
group_gemm_ts_helper.search_cluster_tile_count_k(
cur_tile_coord,
problem_sizes_mnkl,
)
(
cur_k_block_cnt,
cur_group_idx,
) = group_gemm_ts_helper.search_cluster_tile_count_k(
cur_tile_coord,
problem_sizes_mnkl,
)
# Set tensor memory buffer for current tile
acc_buf_idx = tile_sched.num_tiles_executed % self.num_acc_stage
@@ -1022,7 +1025,7 @@ class GroupedGemmKernel:
mma_rd_ab_full_phase = (
(num_prev_k_blk + mma_rd_k_block) // self.num_ab_stage % 2
)
peek_ab_full_status = cute.arch.conditional_mbarrier_try_wait(
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,
@@ -1047,7 +1050,7 @@ class GroupedGemmKernel:
#
# Mma mainloop
#
for k_block in cutlass.range_dynamic(0, cur_k_block_cnt, 1, unroll=1):
for k_block in range(cur_k_block_cnt):
mma_rd_k_block_next = cutlass.Int32(k_block + 1)
smem_rd_buffer_next = (
num_prev_k_blk + mma_rd_k_block_next
@@ -1066,7 +1069,7 @@ class GroupedGemmKernel:
# tCtAcc += tCrA * tCrB
num_kphases = cute.size(tCrA, mode=[2])
for kphase_idx in range(num_kphases):
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
kphase_coord = (None, None, kphase_idx, smem_rd_buffer)
cute.gemm(
@@ -1092,7 +1095,7 @@ class GroupedGemmKernel:
mma_rd_k_block_next < cur_k_block_cnt and is_leader_cta
)
peek_ab_full_status = cute.arch.conditional_mbarrier_try_wait(
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,
@@ -1161,19 +1164,23 @@ class GroupedGemmKernel:
#
# Partition for epilogue
#
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)
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(tma_atom_c, tCgC, epi_tile, sC)
)
(
tma_atom_c,
bSG_sC,
bSG_gC_partitioned,
) = self.epilog_gmem_copy_and_partition(tma_atom_c, tCgC, epi_tile, sC)
#
# Persistent tile scheduling loop
@@ -1270,7 +1277,7 @@ class GroupedGemmKernel:
#
subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt
for subtile_idx in cutlass.range_dynamic(subtile_cnt):
for subtile_idx in range(subtile_cnt):
#
# Load accumulator from tensor memory buffer to register
#
@@ -1493,11 +1500,11 @@ class GroupedGemmKernel:
# (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE)
tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL)
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL)
gC_mnl_epi = cute.flat_divide(
gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
)
# (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL)
# (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(
@@ -1569,14 +1576,14 @@ class GroupedGemmKernel:
- tCgC: The destination global memory tensor partitioned for the TMA operation.
:rtype: tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]
"""
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL)
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL)
gC_epi = cute.flat_divide(
gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
)
sC_for_tma_partition = cute.group_modes(sC, 0, 2)
gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2)
# ((ATOM_V, REST_V), EPI_M, EPI_N)
# ((ATOM_V, REST_V), EPI_M, EPI_N, loopM, loopN, loopL)
# ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL)
bSG_sC, bSG_gC = cpasync.tma_partition(
tma_atom_c,
0,
@@ -1595,7 +1602,7 @@ class GroupedGemmKernel:
epi_tile: cute.Tile,
c_dtype: type[cutlass.Numeric],
c_layout: utils.LayoutEnum,
num_smem_capacity: int,
smem_capacity: int,
occupancy: int,
) -> tuple[int, int, int]:
"""Computes the number of stages for accumulator, A/B operands, and epilogue based on heuristics.
@@ -1614,8 +1621,8 @@ class GroupedGemmKernel:
:type c_dtype: type[cutlass.Numeric]
:param c_layout: Layout enum of operand C in global memory.
:type c_layout: utils.LayoutEnum
:param num_smem_capacity: Total available shared memory capacity in bytes.
:type num_smem_capacity: int
:param smem_capacity: Total available shared memory capacity in bytes.
:type smem_capacity: int
:param occupancy: Target number of CTAs per SM (occupancy).
:type occupancy: int
@@ -1658,7 +1665,7 @@ class GroupedGemmKernel:
# Subtract reserved bytes and initial epilogue bytes
# Divide remaining by bytes needed per A/B stage
num_ab_stage = (
num_smem_capacity // occupancy
smem_capacity // occupancy
- GroupedGemmKernel.reserved_smem_bytes
- epi_bytes
) // ab_bytes_per_stage
@@ -1667,7 +1674,7 @@ class GroupedGemmKernel:
# Calculate remaining smem after allocating for A/B stages and reserved bytes
# Add remaining unused smem to epilogue
remaining_smem = (
num_smem_capacity
smem_capacity
- occupancy * ab_bytes_per_stage * num_ab_stage
- occupancy * (GroupedGemmKernel.reserved_smem_bytes + epi_bytes)
)
@@ -1775,20 +1782,6 @@ class GroupedGemmKernel:
epi_bytes = cute.size_in_bytes(c_dtype, epi_smem_layout_staged)
return ab_bytes + epi_bytes
@staticmethod
def _get_tma_atom_kind(atom_sm_cnt: int, mcast: bool):
"""Select the appropriate TMA copy atom based on the number of SMs and the multicast flag."""
if atom_sm_cnt == 2 and mcast:
return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.TWO)
elif atom_sm_cnt == 2 and not mcast:
return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.TWO)
elif atom_sm_cnt == 1 and mcast:
return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.ONE)
elif atom_sm_cnt == 1 and not mcast:
return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE)
raise ValueError(f"Invalid atom_sm_cnt: {atom_sm_cnt} and {mcast}")
@staticmethod
def _compute_num_tmem_alloc_cols(
tiled_mma: cute.TiledMma,
@@ -1909,8 +1902,6 @@ def run_grouped_gemm(
if not torch.cuda.is_available():
raise RuntimeError("GPU is required to run this example!")
torch.manual_seed(2025)
# Create tensor and return the pointer, tensor, and stride
def create_tensor_and_stride(
l: int,
@@ -1920,42 +1911,17 @@ def run_grouped_gemm(
dtype: type[cutlass.Numeric],
is_dynamic_layout: bool = True,
) -> tuple[int, torch.Tensor, cute.Tensor, torch.Tensor, tuple[int, int]]:
# is_mode0_major: (l, mode1, mode0) -> (mode0, mode1, l)
# else: (l, mode0, mode1) -> (mode0, mode1, l)
shape = (l, mode1, mode0) if is_mode0_major else (l, mode0, mode1)
permute_order = (2, 1, 0) if is_mode0_major else (1, 2, 0)
# omit stride for L mode as it is always 1 for grouped GEMM
strides = (1, mode0) if is_mode0_major else (mode1, 1)
assert dtype in {cutlass.Float16, cutlass.BFloat16, cutlass.Float32}
is_unsigned = False
torch_dtype = cutlass_torch.dtype(dtype)
torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor(
shape,
torch_dtype,
permute_order=permute_order,
init_type=cutlass_torch.TensorInitType.RANDOM,
init_config=cutlass_torch.RandomInitConfig(
min_val=0 if is_unsigned else -2, max_val=4 if is_unsigned else 2
),
torch_tensor_cpu = cutlass_torch.matrix(l, mode0, mode1, is_mode0_major, dtype)
cute_tensor, torch_tensor = cutlass_torch.cute_tensor_like(
torch_tensor_cpu, dtype, is_dynamic_layout, assumed_align=16
)
torch_tensor = torch_tensor_cpu.cuda()
f32_torch_tensor = torch_tensor_cpu.to(dtype=torch.float32)
cute_tensor = from_dlpack(torch_tensor, assumed_align=16)
if is_dynamic_layout:
cute_tensor = cute_tensor.mark_layout_dynamic(
leading_dim=(0 if is_mode0_major else 1)
)
cute_tensor = cutlass_torch.convert_cute_tensor(
f32_torch_tensor,
return (
torch_tensor.data_ptr(),
torch_tensor,
cute_tensor,
dtype,
is_dynamic_layout=is_dynamic_layout,
torch_tensor_cpu,
torch_tensor.stride()[:-1],
)
# Get pointer of the tensor
ptr = torch_tensor.data_ptr()
return ptr, torch_tensor, cute_tensor, f32_torch_tensor, strides
# iterate all groups and create tensors for each group
torch_fp32_tensors_abc = []
@@ -1964,15 +1930,27 @@ def run_grouped_gemm(
strides_abc = []
ptrs_abc = []
for _, (m, n, k, l) in enumerate(problem_sizes_mnkl):
ptr_a, torch_tensor_a, cute_tensor_a, tensor_fp32_a, stride_mk_a = (
create_tensor_and_stride(l, m, k, a_major == "m", ab_dtype)
)
ptr_b, torch_tensor_b, cute_tensor_b, tensor_fp32_b, stride_nk_b = (
create_tensor_and_stride(l, n, k, b_major == "n", ab_dtype)
)
ptr_c, torch_tensor_c, cute_tensor_c, tensor_fp32_c, stride_mn_c = (
create_tensor_and_stride(l, m, n, c_major == "m", c_dtype)
)
(
ptr_a,
torch_tensor_a,
cute_tensor_a,
tensor_fp32_a,
stride_mk_a,
) = create_tensor_and_stride(l, m, k, a_major == "m", ab_dtype)
(
ptr_b,
torch_tensor_b,
cute_tensor_b,
tensor_fp32_b,
stride_nk_b,
) = create_tensor_and_stride(l, n, k, b_major == "n", ab_dtype)
(
ptr_c,
torch_tensor_c,
cute_tensor_c,
tensor_fp32_c,
stride_mn_c,
) = create_tensor_and_stride(l, m, n, c_major == "m", c_dtype)
ptrs_abc.append([ptr_a, ptr_b, ptr_c])
torch_tensors_abc.append([torch_tensor_a, torch_tensor_b, torch_tensor_c])
torch_fp32_tensors_abc.append([tensor_fp32_a, tensor_fp32_b, tensor_fp32_c])
@@ -2005,19 +1983,16 @@ def run_grouped_gemm(
)
# Prepare tensormap buffer for each SM
num_tensormap_buffers = sm_count
tensormap_pytorch_tensor = (
torch.empty(
(
num_tensormap_buffers,
GroupedGemmKernel.num_tensormaps,
GroupedGemmKernel.bytes_per_tensormap // 8,
),
dtype=torch.int64,
)
.fill_(0)
.cuda()
tensormap_shape = (
num_tensormap_buffers,
GroupedGemmKernel.num_tensormaps,
GroupedGemmKernel.bytes_per_tensormap // 8,
)
tensor_of_tensormap, tensor_of_tensormap_torch = cutlass_torch.cute_tensor_like(
torch.empty(tensormap_shape, dtype=torch.int64),
cutlass.Int64,
is_dynamic_layout=False,
)
tensormap_cute_tensor = from_dlpack(tensormap_pytorch_tensor, assumed_align=16)
grouped_gemm = GroupedGemmKernel(
acc_dtype,
@@ -2027,23 +2002,30 @@ def run_grouped_gemm(
tensormap_update_mode,
)
# Convert integer list to torch tensor and cute tensor
def convert_list_to_tensor(l, dtype) -> tuple[torch.Tensor, cute.Tensor]:
torch_tensor = torch.tensor(l, dtype=dtype).cuda()
cute_tensor = from_dlpack(torch_tensor, assumed_align=16)
return torch_tensor, cute_tensor
# layout (num_groups, 4):(4, 1)
problem_sizes_mnkl_torch_tensor, problem_sizes_mnkl_cute_tensor = (
convert_list_to_tensor(problem_sizes_mnkl, torch.int32)
(
tensor_of_dim_size_mnkl,
tensor_of_dim_size_mnkl_torch,
) = cutlass_torch.cute_tensor_like(
torch.tensor(problem_sizes_mnkl, dtype=torch.int32),
cutlass.Int32,
is_dynamic_layout=False,
assumed_align=16,
)
# layout (num_groups, 3, 2):(6, 2, 1)
strides_abc_torch_tensor, strides_abc_cute_tensor = convert_list_to_tensor(
strides_abc, torch.int32
tensor_of_strides_abc, tensor_of_strides_abc_torch = cutlass_torch.cute_tensor_like(
torch.tensor(strides_abc, dtype=torch.int32),
cutlass.Int32,
is_dynamic_layout=False,
assumed_align=16,
)
# layout (num_groups,3):(3, 1)
ptrs_abc_torch_tensor, ptrs_abc_cute_tensor = convert_list_to_tensor(
ptrs_abc, torch.int64
tensor_of_ptrs_abc, tensor_of_ptrs_abc_torch = cutlass_torch.cute_tensor_like(
torch.tensor(ptrs_abc, dtype=torch.int64),
cutlass.Int64,
is_dynamic_layout=False,
assumed_align=16,
)
# Compute total number of cluster tiles we need to compute for given grouped GEMM problem
@@ -2077,10 +2059,9 @@ def run_grouped_gemm(
problem_sizes_mnkl, cluster_tile_shape_mn
)
# Get current CUDA stream from PyTorch
torch_stream = torch.cuda.current_stream()
# Get the raw stream pointer as a CUstream
current_stream = cuda.CUstream(torch_stream.cuda_stream)
# Initialize Stream
current_stream = cutlass_torch.default_stream()
# Compile grouped GEMM kernel
compiled_grouped_gemm = cute.compile(
grouped_gemm,
@@ -2088,11 +2069,11 @@ def run_grouped_gemm(
initial_cute_tensors_abc[1],
initial_cute_tensors_abc[2],
num_groups,
problem_sizes_mnkl_cute_tensor,
strides_abc_cute_tensor,
ptrs_abc_cute_tensor,
tensor_of_dim_size_mnkl,
tensor_of_strides_abc,
tensor_of_ptrs_abc,
total_num_clusters,
tensormap_cute_tensor,
tensor_of_tensormap,
max_active_clusters,
current_stream,
)
@@ -2104,10 +2085,10 @@ def run_grouped_gemm(
initial_cute_tensors_abc[0],
initial_cute_tensors_abc[1],
initial_cute_tensors_abc[2],
problem_sizes_mnkl_cute_tensor,
strides_abc_cute_tensor,
ptrs_abc_cute_tensor,
tensormap_cute_tensor,
tensor_of_dim_size_mnkl,
tensor_of_strides_abc,
tensor_of_ptrs_abc,
tensor_of_tensormap,
current_stream,
)
# Execution
@@ -2116,28 +2097,27 @@ def run_grouped_gemm(
initial_cute_tensors_abc[0],
initial_cute_tensors_abc[1],
initial_cute_tensors_abc[2],
problem_sizes_mnkl_cute_tensor,
strides_abc_cute_tensor,
ptrs_abc_cute_tensor,
tensormap_cute_tensor,
tensor_of_dim_size_mnkl,
tensor_of_strides_abc,
tensor_of_ptrs_abc,
tensor_of_tensormap,
current_stream,
)
torch.cuda.synchronize()
# Compute reference result
if not skip_ref_check:
refs = []
for a, b, _ in torch_fp32_tensors_abc:
ref = (torch.einsum("mkl,nkl->mnl", a, b)).cpu()
refs.append(ref)
for i, ((_, _, c), ref) in enumerate(zip(torch_tensors_abc, refs)):
for i, (a, b, c) in enumerate(torch_tensors_abc):
ref = torch.einsum(
"mkl,nkl->mnl",
a.cpu().to(dtype=torch.float32),
b.cpu().to(dtype=torch.float32),
)
print(f"checking group {i}")
if c_dtype == cutlass.Float32:
ref_c = ref
else:
ref_c = ref.to(cutlass_torch.dtype(c_dtype))
torch.testing.assert_close(
c.cpu(),
ref_c,
ref.to(cutlass_torch.dtype(c_dtype)),
atol=tolerance,
rtol=1e-05,
)
@@ -2266,6 +2246,8 @@ if __name__ == "__main__":
else:
tensormap_update_mode = utils.TensorMapUpdateMode.SMEM
torch.manual_seed(2025)
run_grouped_gemm(
args.num_groups,
args.problem_sizes_mnkl,

3619
examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py Normal file
View File

File diff suppressed because it is too large Load Diff

397
examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd_reference.py Normal file
View File

@@ -0,0 +1,397 @@
# 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 torch.nn.functional as F
def ssd_reference_fp32_all(x, a, delta, B, C, Y_out, Fstate_out, D, has_d, d_has_hdim):
"""
Rearrange tensor dimensions from cuda layout to reference layout, then directly call TriDao's ssd implementation
Arguments:
X/x: (D, L, C, H, B):(C*L, 1, L, D*C*L, H*D*C*L)
A/delta: (L, C, H, B):(1, L, C*L, H*C*L)
a: (H):(1)
B/C: (L, N, C, G, B):(1, C*L, L, N*C*L, G*N*C*L)
D: (1, H):(0, 1) or (D, H):(1, D)
has_d: bool
d_has_hdim: bool
Return:
Y_out: (L, D, C, H, B):(1, C*L, L, D*C*L, H*D*C*L)
Fstate_out: (D, N, H, B):(N, 1, D*N, H*D*N)
"""
assert x.dtype == a.dtype == delta.dtype == B.dtype == C.dtype
A = delta * a.view(1, 1, -1, 1)
X = x * delta.unsqueeze(0)
# Rearrange to match cutlass layout to tridao's layout
block_len = A.shape[0]
initial_states = None
# A: l c h b-> b c l h
A = A.permute(3, 1, 0, 2)
# X: p l c h b -> b c l h p
X = X.permute(4, 2, 1, 3, 0)
# B: l n c g b -> b c l g n
B = B.permute(4, 2, 0, 3, 1)
# C: l n c g b -> b c l g n
C = C.permute(4, 2, 0, 3, 1)
# X/A/B/C: b c l ... -> b (c l) ...
X, A, B, C = [x.reshape(x.shape[0], -1, *x.shape[3:]) for x in (X, A, B, C)]
# Ngroup (g to h) mapping
B_val, CL_val, G_val, N_val = B.shape
H_val = X.shape[2]
ngroup_ratio = H_val // G_val
# B/C: (B, CL, H, N)
h_to_g_mapping = torch.arange(H_val, device=B.device) // ngroup_ratio
B = B.gather(2, h_to_g_mapping.view(1, 1, -1, 1).expand(B_val, CL_val, -1, N_val))
C = C.gather(2, h_to_g_mapping.view(1, 1, -1, 1).expand(B_val, CL_val, -1, N_val))
###################################################################
# Call reference implementation from Tri Dao ssd_minimal_discrete
Y, final_state = ssd_minimal_discrete_fp32_all(
X, A, B, C, block_len, initial_states
)
###################################################################
if has_d:
D_val = Y.shape[3]
if not d_has_hdim:
D = D.expand(D_val, -1)
Y = Y + torch.einsum("bchp,ph->bchp", X, D)
# Rearrange to match tridao's layout to cutlass layout
# Y: b (c l) h p -> b c l h p
Y = Y.reshape(Y.shape[0], -1, block_len, Y.shape[2], Y.shape[3])
# Y: b c l h p -> l p c h b
Y = Y.permute(2, 4, 1, 3, 0)
# Fstate_out: b h p n -> p n h b
Fstate_out.copy_(final_state.permute(2, 3, 1, 0))
Y_out.copy_(Y)
return
def ssd_reference_lowprecision_intermediates(
x, a, delta, B, C, Y_out, Fstate_out, intermediate_dtype, D, has_d, d_has_hdim
):
"""
Rearrange tensor dimensions from cuda layout to reference layout, then call a reduced intermediate dtype version of ssd implementation
Arguments:
X/x: (D, L, C, H, B):(C*L, 1, L, D*C*L, H*D*C*L)
A/delta: (L, C, H, B):(1, L, C*L, H*C*L)
a: (H):(1)
B/C: (L, N, C, G, B):(1, C*L, L, N*C*L, G*N*C*L)
intermediate_dtype: input and intermediate data type
D: (1, H):(0, 1) or (D, H):(1, D)
has_d: bool
d_has_hdim: bool
Return:
Y_out: (L, D, C, H, B):(1, C*L, L, D*C*L, H*D*C*L)
Fstate_out: (D, N, H, B):(N, 1, D*N, H*D*N)
"""
assert x.dtype == a.dtype == delta.dtype == B.dtype == C.dtype
A = delta * a.view(1, 1, -1, 1)
# Rearrange to match cutlass layout to tridao's layout
block_len = A.shape[0]
initial_states = None
# A: l c h b-> b c l h
A = A.permute(3, 1, 0, 2)
# delta: l c h b-> b c l h
delta = delta.permute(3, 1, 0, 2)
# x: p l c h b -> b c l h p
x = x.permute(4, 2, 1, 3, 0)
# B: l n c g b -> b c l g n
B = B.permute(4, 2, 0, 3, 1)
# C: l n c g b -> b c l g n
C = C.permute(4, 2, 0, 3, 1)
# x/A/delta/B/C: b c l ... -> b (c l) ...
x, A, delta, B, C = [
tensor.reshape(tensor.shape[0], -1, *tensor.shape[3:])
for tensor in (x, A, delta, B, C)
]
# Ngroup (g to h) mapping
B_val, CL_val, G_val, N_val = B.shape
H_val = x.shape[2]
ngroup_ratio = H_val // G_val
# B/C: (B, CL, H, N)
h_to_g_mapping = torch.arange(H_val, device=B.device) // ngroup_ratio
B = B.gather(2, h_to_g_mapping.view(1, 1, -1, 1).expand(B_val, CL_val, -1, N_val))
C = C.gather(2, h_to_g_mapping.view(1, 1, -1, 1).expand(B_val, CL_val, -1, N_val))
# Type convert input tensors to input dtype (same as intermediate dtype)
x = x.to(intermediate_dtype).to(torch.float32)
A = A.to(intermediate_dtype).to(torch.float32)
delta = delta.to(intermediate_dtype).to(torch.float32)
B = B.to(intermediate_dtype).to(torch.float32)
C = C.to(intermediate_dtype).to(torch.float32)
#########################################################################
# Call reference implementation ssd_minimal_discrete_bf16_intermediates
Y, final_state = ssd_minimal_discrete_lowprecision_intermediates(
x, A, delta, B, C, block_len, intermediate_dtype, initial_states
)
#########################################################################
if has_d:
D = D.to(intermediate_dtype).to(torch.float32)
D_val = Y.shape[3]
if not d_has_hdim:
D = D.expand(D_val, -1)
Y = Y + torch.einsum("bchp,ph->bchp", x, D)
# Type convert output tensors to output dtype (same as intermediate dtype)
Y = Y.to(intermediate_dtype).to(torch.float32)
final_state = final_state.to(intermediate_dtype).to(torch.float32)
# Rearrange to match tridao's layout to cutlass layout
# Y: b (c l) h p -> b c l h p
Y = Y.reshape(Y.shape[0], -1, block_len, Y.shape[2], Y.shape[3])
# Y: b c l h p -> l p c h b
Y = Y.permute(2, 4, 1, 3, 0)
# Fstate_out: b h p n -> p n h b
Fstate_out.copy_(final_state.permute(2, 3, 1, 0))
Y_out.copy_(Y)
return
def analyze_relative_diffs(actual, expected):
"""
Print statistics of relative differences between actual and expected tensors
"""
# Calculate relative differences
abs_diff = (actual - expected).abs()
rel_diff = abs_diff / (torch.maximum(expected.abs(), actual.abs()) + 0.00001)
total_elements = rel_diff.numel()
# Handle special cases first
nan_mask = torch.isnan(rel_diff)
inf_mask = torch.isinf(rel_diff)
nan_count = nan_mask.sum().item()
inf_count = inf_mask.sum().item()
# Find position and value of maximum relative difference
max_rel_diff = (
rel_diff[~nan_mask & ~inf_mask].max()
if (~nan_mask & ~inf_mask).any()
else float("nan")
)
max_rel_diff_pos = (
rel_diff[~nan_mask & ~inf_mask].argmax()
if (~nan_mask & ~inf_mask).any()
else -1
)
# Print max relative difference info
print(f"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]}")
print(f"Expected value: {expected.flatten()[max_rel_diff_pos]}")
print(f"NaN values: {nan_count} ({100.0 * nan_count / total_elements:.2f}%)")
print(f"Inf values: {inf_count} ({100.0 * inf_count / total_elements:.2f}%)\n")
# Check different rtol thresholds
rtol_levels = [1e-5, 1e-4, 1e-3, 1e-2, 5e-02, 1e-01]
for i, rtol in enumerate(rtol_levels):
if i == 0:
mask = rel_diff <= rtol
else:
mask = (rel_diff <= rtol) & (rel_diff > rtol_levels[i - 1])
count = mask.sum().item()
percentage = (count / total_elements) * 100
if i == 0:
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}%)"
)
# Print elements exceeding the largest rtol
mask = rel_diff > rtol_levels[-1]
count = mask.sum().item()
percentage = (count / total_elements) * 100
print(f"Elements with rtol > {rtol_levels[-1]:.0e}: {count} ({percentage:.2f}%)\n")
def segsum(x):
"""
More stable segment sum calculation.
x: b h c l
"""
T = x.size(-1)
# x: b h c l -> b h c l l
x = x.unsqueeze(-1).expand(*x.shape, T)
mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=-1)
x = x.masked_fill(~mask, 0)
x_segsum = torch.cumsum(x, dim=-2)
mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=0)
x_segsum = x_segsum.masked_fill(~mask, -torch.inf)
return x_segsum
def ssd_minimal_discrete_fp32_all(X, A, B, C, block_len, initial_states=None):
"""
This is same with https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/ssd_minimal.py
(all accumulation and intermediate results in fp32)
Arguments:
X: (batch(B), length(C*L), n_heads(H), d_head(D))
A: (batch(B), length(C*L), n_heads(H))
B: (batch(B), length(C*L), n_heads(H), d_state(N))
C: (batch(B), length(C*L), n_heads(H), d_state(N))
Return:
Y: (batch(B), length(C*L), n_heads(H), d_head(D))
final_state: (B, H, D, N)
"""
assert X.dtype == A.dtype == B.dtype == C.dtype
assert X.shape[1] % block_len == 0
# Rearrange into blocks/chunks
# X/A/B/C:b (c l) ... -> b c l ...
X, A, B, C = [
x.reshape(x.shape[0], -1, block_len, *x.shape[2:]) for x in (X, A, B, C)
]
# A: b c l h -> b h c l
A = A.permute(0, 3, 1, 2)
# A_cumsum: (B, H, C, L)
A_cumsum = torch.cumsum(A, dim=-1)
# 1. Compute the output for each intra-chunk (diagonal blocks)
segsum_A = segsum(A)
L = torch.exp(segsum_A)
Y_diag = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", C, B, L, X)
# 2. Compute the state for each intra-chunk
# (right term of low-rank factorization of off-diagonal blocks; B terms)
decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum))
states = torch.einsum("bclhn,bhcl,bclhp->bchpn", B, decay_states, X)
# 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
# (middle term of factorization of off-diag blocks; A terms)
if initial_states is None:
initial_states = torch.zeros_like(states[:, :1])
states = torch.cat([initial_states, states], dim=1)
decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
states, final_state = new_states[:, :-1], new_states[:, -1]
# 4. Compute state -> output conversion per chunk
# (left term of low-rank factorization of off-diagonal blocks; C terms)
state_decay_out = torch.exp(A_cumsum)
Y_off = torch.einsum("bclhn,bchpn,bhcl->bclhp", C, states, state_decay_out)
# Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
# Y: b c l h p -> b (c l) h p
Y = (Y_diag + Y_off).reshape(Y_diag.shape[0], -1, Y_diag.shape[3], Y_diag.shape[4])
return Y, final_state
def ssd_minimal_discrete_lowprecision_intermediates(
X, A, delta, B, C, block_len, intermediate_dtype, initial_states=None
):
"""
This is adjusted from ssd_minimal_discrete_fp32_all, with exceptions:
1. accumulation in fp32 but intermediates Q/b_tmem/P are in intermediate_dtype
2. delta is not pre-multiplied with X, delta was applied to generate Q/b_tmem to match GPU implementation
Arguments:
X: (batch(B), length(C*L), n_heads(H), d_head(D))
A: (batch(B), length(C*L), n_heads(H))
delta: (batch(B), length(C*L), n_heads(H))
B: (batch(B), length(C*L), n_heads(H), d_state(N))
C: (batch(B), length(C*L), n_heads(H), d_state(N))
Return:
Y: (batch(B), length(C*L), n_heads(H), d_head(D))
final_state: (B, H, D, N)
"""
assert X.dtype == A.dtype == B.dtype == C.dtype
assert X.shape[1] % block_len == 0
# Rearrange into blocks/chunks
# X/A/delta/B/C: b (c l) ... -> b c l ...
X, A, delta, B, C = [
x.reshape(x.shape[0], -1, block_len, *x.shape[2:]) for x in (X, A, delta, B, C)
]
# A: b c l h -> b h c l
A = A.permute(0, 3, 1, 2)
# delta: b c l h -> b h c l
delta = delta.permute(0, 3, 1, 2)
# A_cumsum: (B, H, C, L)
A_cumsum = torch.cumsum(A, dim=-1)
# 1. Compute the output for each intra-chunk (diagonal blocks)
segsum_A = segsum(A)
L = torch.exp(segsum_A)
intra_acc_0 = torch.einsum("bclhn,bcshn->bclhs", C, B)
Q = torch.einsum("bclhs,bhcls,bhcs->bclhs", intra_acc_0, L, delta)
Y_diag = torch.einsum(
"bclhs,bcshp->bclhp", Q.to(intermediate_dtype).to(torch.float32), X
)
# 2. Compute the state for each intra-chunk
# (right term of low-rank factorization of off-diagonal blocks; B terms)
decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum))
b_tmem = torch.einsum("bclhn,bhcl,bhcl->bclhn", B, decay_states, delta)
states = torch.einsum(
"bclhn,bclhp->bchpn", b_tmem.to(intermediate_dtype).to(torch.float32), X
)
# 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
# (middle term of factorization of off-diag blocks; A terms)
if initial_states is None:
initial_states = torch.zeros_like(states[:, :1])
states = torch.cat([initial_states, states], dim=1)
decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
states, final_state = new_states[:, :-1], new_states[:, -1]
final_state = final_state
# 4. Compute state -> output conversion per chunk
# (left term of low-rank factorization of off-diagonal blocks; C terms)
state_decay_out = torch.exp(A_cumsum)
Y_off_tmp = torch.einsum(
"bclhn,bchpn->bclhp", C, states.to(intermediate_dtype).to(torch.float32)
)
Y_off = torch.einsum("bclhp,bhcl->bclhp", Y_off_tmp, state_decay_out)
# Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
# Y: b c l h p -> b (c l) h p
Y = (Y_diag + Y_off).reshape(
Y_diag.shape[0], -1, Y_diag.shape[3], Y_diag.shape[4]
) # b (c l) h p
return Y, final_state

200
examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd_tile_scheduler.py Normal file
View File

@@ -0,0 +1,200 @@
# 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 typing import Tuple
from cutlass.cutlass_dsl import (
Boolean,
Integer,
Int32,
min,
extract_mlir_values,
new_from_mlir_values,
dsl_user_op,
)
from cutlass._mlir import ir
import cutlass.cute as cute
from cutlass.utils import WorkTileInfo
class Mamba2SSDTileSchedulerParams:
def __init__(
self,
problem_shape_ntiles: int,
eh: int,
ngroup_ratio: int,
*,
loc=None,
ip=None,
):
self.problem_shape_ntiles = problem_shape_ntiles
self.eh = eh
self.ngroup_ratio = ngroup_ratio
self._loc = loc
def __extract_mlir_values__(self):
values, self._values_pos = [], []
for obj in [self.problem_shape_ntiles, self.eh, self.ngroup_ratio]:
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_ntiles, self.eh, self.ngroup_ratio], self._values_pos
):
obj_list.append(new_from_mlir_values(obj, values[:n_items]))
values = values[n_items:]
return Mamba2SSDTileSchedulerParams(*(tuple(obj_list)), loc=self._loc)
@dsl_user_op
def get_grid_shape(
self, max_active_clusters: Int32, *, loc=None, ip=None
) -> Tuple[Integer, Integer, Integer]:
return (min(self.problem_shape_ntiles, max_active_clusters), 1, 1)
class Mamba2SSDTileScheduler:
def __init__(
self,
params: Mamba2SSDTileSchedulerParams,
num_persistent_ctas: Int32,
current_work_linear_idx: Int32,
num_tiles_executed: Int32,
):
self.params = params
self.num_persistent_ctas = num_persistent_ctas
self._current_work_linear_idx = current_work_linear_idx
self._num_tiles_executed = num_tiles_executed
def __extract_mlir_values__(self) -> list[ir.Value]:
values = extract_mlir_values(self.num_persistent_ctas)
values.extend(extract_mlir_values(self._current_work_linear_idx))
values.extend(extract_mlir_values(self._num_tiles_executed))
return values
def __new_from_mlir_values__(
self, values: list[ir.Value]
) -> "Mamba2SSDTileScheduler":
assert len(values) == 3
new_num_persistent_ctas = new_from_mlir_values(
self.num_persistent_ctas, [values[0]]
)
new_current_work_linear_idx = new_from_mlir_values(
self._current_work_linear_idx, [values[1]]
)
new_num_tiles_executed = new_from_mlir_values(
self._num_tiles_executed, [values[2]]
)
return Mamba2SSDTileScheduler(
self.params,
new_num_persistent_ctas,
new_current_work_linear_idx,
new_num_tiles_executed,
)
# called by host
@dsl_user_op
@staticmethod
def create(
params: Mamba2SSDTileSchedulerParams,
block_idx: Tuple[Integer, Integer, Integer],
grid_dim: Tuple[Integer, Integer, Integer],
*,
loc=None,
ip=None,
):
params = params
# Calculate the number of persistent clusters by dividing the total grid size
# by the number of CTAs per cluster
num_persistent_ctas = Int32(cute.size(grid_dim, loc=loc, ip=ip))
bidx, bidy, bidz = block_idx
# Initialize workload index equals to the cluster index in the grid
current_work_linear_idx = Int32(bidx)
# Initialize number of tiles executed to zero
num_tiles_executed = Int32(0)
return Mamba2SSDTileScheduler(
params,
num_persistent_ctas,
current_work_linear_idx,
num_tiles_executed,
)
# called by host
@staticmethod
def get_grid_shape(
params: Mamba2SSDTileSchedulerParams,
max_active_clusters: Int32,
*,
loc=None,
ip=None,
) -> Tuple[Integer, Integer, Integer]:
return params.get_grid_shape(max_active_clusters, loc=loc, ip=ip)
# private method
def _get_current_work_for_linear_idx(
self, current_work_linear_idx: Int32, *, loc=None, ip=None
) -> WorkTileInfo:
is_valid = current_work_linear_idx < cute.size(
self.params.problem_shape_ntiles, loc=loc, ip=ip
)
eh_idx = current_work_linear_idx % self.params.eh
b_idx = current_work_linear_idx // self.params.eh
g_idx = eh_idx // self.params.ngroup_ratio
# cur_tile_coord is (b_idx, eh_idx, g_idx)
cur_tile_coord = tuple(Int32(x) for x in (b_idx, eh_idx, g_idx))
return WorkTileInfo(cur_tile_coord, is_valid)
@dsl_user_op
def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
return self._get_current_work_for_linear_idx(
self._current_work_linear_idx, loc=loc, ip=ip
)
@dsl_user_op
def initial_work_tile_info(self, *, loc=None, ip=None) -> WorkTileInfo:
return self.get_current_work(loc=loc, ip=ip)
@dsl_user_op
def advance_to_next_work(self, *, advance_count: int = 1, loc=None, ip=None):
self._current_work_linear_idx += Int32(advance_count) * Int32(
self.num_persistent_ctas
)
self._num_tiles_executed += Int32(1)
@property
def num_tiles_executed(self) -> Int32:
return self._num_tiles_executed