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

v4.3 tag release update. (#2789)

Browse Source
This commit is contained in:
Junkai-Wu
2025-11-20 20:49:44 -05:00
committed by GitHub
parent 406e078b29
commit 8cd5bef43a
225 changed files with 23229 additions and 2813 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py
+12 -14
View File
@@ -39,6 +39,7 @@ from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
import cutlass.utils.blackwell_helpers as sm100_utils
import math
@@ -220,22 +221,18 @@ class BlockwiseGemmKernel:
self.num_regs_epilogue_warps = 216
self.num_regs_acc_update_warps = 216
# Set barrier for cta sync, epilogue sync and tmem ptr sync
self.cta_sync_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=self.threads_per_cta,
)
# Set barrier for epilogue sync and tmem ptr sync
self.epilog_sync_barrier = pipeline.NamedBarrier(
barrier_id=2,
barrier_id=1,
num_threads=32 * len(self.epilog_warp_id),
)
self.tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=3,
barrier_id=2,
num_threads=32
* len((self.mma_warp_id, *self.epilog_warp_id, *self.acc_update_warp_id)),
)
self.sched_sync_barrier = pipeline.NamedBarrier(
barrier_id=4,
barrier_id=3,
num_threads=self.threads_per_warp,
)
self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
@@ -703,6 +700,7 @@ class BlockwiseGemmKernel:
consumer_group=ab_pipeline_consumer_group,
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Initialize mainloop scale_pipeline (barrier) and states
@@ -719,6 +717,7 @@ class BlockwiseGemmKernel:
num_stages=self.num_scale_stage,
producer_group=scale_pipeline_producer_group,
consumer_group=scale_pipeline_consumer_group,
defer_sync=True,
)
# Initialize acc_pipeline (barrier) and states
@@ -735,6 +734,7 @@ class BlockwiseGemmKernel:
producer_group=acc_pipeline_producer_group,
consumer_group=acc_pipeline_consumer_group,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Initialize epilogue pipeline (barrier) and states
@@ -751,6 +751,7 @@ class BlockwiseGemmKernel:
num_stages=1,
producer_group=epi_pipeline_producer_group,
consumer_group=epi_pipeline_consumer_group,
defer_sync=True,
)
# Initialize tile info pipeline (barrier) and states
@@ -767,6 +768,7 @@ class BlockwiseGemmKernel:
num_stages=self.num_tile_stage,
producer_group=tile_info_pipeline_producer_group,
consumer_group=tile_info_pipeline_consumer_group,
defer_sync=True,
)
# Tensor memory dealloc barrier init
@@ -779,8 +781,7 @@ class BlockwiseGemmKernel:
)
# Cluster arrive after barrier init
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_arrive_relaxed()
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
#
# Setup smem tensor A/B/C/Scale
@@ -968,10 +969,7 @@ class BlockwiseGemmKernel:
#
# Cluster wait before tensor memory alloc
#
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
else:
self.cta_sync_barrier.arrive_and_wait()
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
#
# Specialized Schedule warp
examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py
+12 -14
View File
@@ -39,6 +39,7 @@ from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
import cutlass.utils.blackwell_helpers as sm100_utils
from cutlass.cute.runtime import from_dlpack
@@ -236,22 +237,18 @@ class BlockwiseContiguousGroupedGemmKernel:
self.num_regs_epilogue_warps = 216
self.num_regs_acc_update_warps = 216
# Set barrier for cta sync, epilogue sync and tmem ptr sync
self.cta_sync_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=self.threads_per_cta,
)
# Set barrier for epilogue sync and tmem ptr sync
self.epilog_sync_barrier = pipeline.NamedBarrier(
barrier_id=2,
barrier_id=1,
num_threads=32 * len(self.epilog_warp_id),
)
self.tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=3,
barrier_id=2,
num_threads=32
* len((self.mma_warp_id, *self.epilog_warp_id, *self.acc_update_warp_id)),
)
self.sched_sync_barrier = pipeline.NamedBarrier(
barrier_id=4,
barrier_id=3,
num_threads=self.threads_per_warp,
)
self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
@@ -724,6 +721,7 @@ class BlockwiseContiguousGroupedGemmKernel:
consumer_group=ab_pipeline_consumer_group,
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Initialize mainloop scale_pipeline (barrier) and states
@@ -740,6 +738,7 @@ class BlockwiseContiguousGroupedGemmKernel:
num_stages=self.num_scale_stage,
producer_group=scale_pipeline_producer_group,
consumer_group=scale_pipeline_consumer_group,
defer_sync=True,
)
# Initialize acc_pipeline (barrier) and states
@@ -756,6 +755,7 @@ class BlockwiseContiguousGroupedGemmKernel:
producer_group=acc_pipeline_producer_group,
consumer_group=acc_pipeline_consumer_group,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Initialize epilogue pipeline (barrier) and states
@@ -772,6 +772,7 @@ class BlockwiseContiguousGroupedGemmKernel:
num_stages=1,
producer_group=epi_pipeline_producer_group,
consumer_group=epi_pipeline_consumer_group,
defer_sync=True,
)
# Initialize tile info pipeline (barrier) and states
@@ -788,6 +789,7 @@ class BlockwiseContiguousGroupedGemmKernel:
num_stages=self.num_tile_stage,
producer_group=tile_info_pipeline_producer_group,
consumer_group=tile_info_pipeline_consumer_group,
defer_sync=True,
)
# Tensor memory dealloc barrier init
@@ -800,8 +802,7 @@ class BlockwiseContiguousGroupedGemmKernel:
)
# Cluster arrive after barrier init
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_arrive_relaxed()
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
#
# Setup smem tensor A/B/C/Scale
@@ -989,10 +990,7 @@ class BlockwiseContiguousGroupedGemmKernel:
#
# Cluster wait before tensor memory alloc
#
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
else:
self.cta_sync_barrier.arrive_and_wait()
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
#
# Specialized Schedule warp
examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py
+12 -14
View File
@@ -39,6 +39,7 @@ from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
import cutlass.utils.blackwell_helpers as sm100_utils
from cutlass.cute.runtime import from_dlpack
@@ -235,22 +236,18 @@ class BlockwiseMaskedGroupedGemmKernel:
self.num_regs_epilogue_warps = 216
self.num_regs_acc_update_warps = 216
# Set barrier id for cta sync, epilogue sync and tmem ptr sync
self.cta_sync_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=self.threads_per_cta,
)
# Set barrier id for epilogue sync and tmem ptr sync
self.epilog_sync_barrier = pipeline.NamedBarrier(
barrier_id=2,
barrier_id=1,
num_threads=32 * len(self.epilog_warp_id),
)
self.tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=3,
barrier_id=2,
num_threads=32
* len((self.mma_warp_id, *self.epilog_warp_id, *self.acc_update_warp_id)),
)
self.sched_sync_barrier = pipeline.NamedBarrier(
barrier_id=4,
barrier_id=3,
num_threads=self.threads_per_warp,
)
self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
@@ -723,6 +720,7 @@ class BlockwiseMaskedGroupedGemmKernel:
consumer_group=ab_pipeline_consumer_group,
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Initialize mainloop scale_pipeline (barrier) and states
@@ -739,6 +737,7 @@ class BlockwiseMaskedGroupedGemmKernel:
num_stages=self.num_scale_stage,
producer_group=scale_pipeline_producer_group,
consumer_group=scale_pipeline_consumer_group,
defer_sync=True,
)
# Initialize acc_pipeline (barrier) and states
@@ -755,6 +754,7 @@ class BlockwiseMaskedGroupedGemmKernel:
producer_group=acc_pipeline_producer_group,
consumer_group=acc_pipeline_consumer_group,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Initialize epilogue pipeline (barrier) and states
@@ -771,6 +771,7 @@ class BlockwiseMaskedGroupedGemmKernel:
num_stages=1,
producer_group=epi_pipeline_producer_group,
consumer_group=epi_pipeline_consumer_group,
defer_sync=True,
)
# Initialize tile info pipeline (barrier) and states
@@ -787,6 +788,7 @@ class BlockwiseMaskedGroupedGemmKernel:
num_stages=self.num_tile_stage,
producer_group=tile_info_pipeline_producer_group,
consumer_group=tile_info_pipeline_consumer_group,
defer_sync=True,
)
# Tensor memory dealloc barrier init
@@ -799,8 +801,7 @@ class BlockwiseMaskedGroupedGemmKernel:
)
# Cluster arrive after barrier init
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_arrive_relaxed()
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
#
# Setup smem tensor A/B/C/Scale
@@ -988,10 +989,7 @@ class BlockwiseMaskedGroupedGemmKernel:
#
# Cluster wait before tensor memory alloc
#
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
else:
self.cta_sync_barrier.arrive_and_wait()
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
#
# Specialized Schedule warp
examples/python/CuTeDSL/blackwell/dense_blockscaled_gemm_persistent.py
+96 -34
View File
@@ -38,6 +38,7 @@ from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.cute.runtime import from_dlpack
@@ -208,17 +209,13 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
self.threads_per_cta = 32 * len(
(self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
)
# Set barrier id for cta sync, epilogue sync and tmem ptr sync
self.cta_sync_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=self.threads_per_cta,
)
# Set barrier id for epilogue sync and tmem ptr sync
self.epilog_sync_barrier = pipeline.NamedBarrier(
barrier_id=2,
barrier_id=1,
num_threads=32 * len(self.epilog_warp_id),
)
self.tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=3,
barrier_id=2,
num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
)
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
@@ -288,6 +285,11 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
self.mma_tiler[1],
self.mma_tiler[2],
)
self.cta_tile_shape_mnk_sfb = (
self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape),
self.mma_tiler_sfb[1],
self.mma_tiler_sfb[2],
)
# Compute cluster layout
self.cluster_layout_vmnk = cute.tiled_divide(
@@ -314,6 +316,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
self.c_layout,
self.c_dtype,
)
self.epi_tile_n = cute.size(self.epi_tile[1])
# 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_c_stage = self._compute_stages(
@@ -362,6 +365,19 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
self.num_c_stage,
)
# Overlap and double buffer accumulator when num_acc_stage == 1 for cta_tile_n = 256 case
self.overlapping_accum = self.num_acc_stage == 1
# Compute number of TMEM columns for SFA/SFB/Accumulator
sf_atom_mn = 32
self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k
self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k
self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols
self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] * self.num_acc_stage if not self.overlapping_accum else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols
# Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue
self.iter_acc_early_release_in_epilogue = self.num_sf_tmem_cols // self.epi_tile_n
@cute.jit
def __call__(
self,
@@ -640,6 +656,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
stream=stream,
min_blocks_per_mp=1,
)
return
@@ -726,6 +743,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
consumer_group=ab_pipeline_consumer_group,
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Initialize acc_pipeline (barrier) and states
@@ -742,6 +760,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
producer_group=acc_pipeline_producer_group,
consumer_group=acc_pipeline_consumer_group,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Tensor memory dealloc barrier init
@@ -754,8 +773,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
)
# Cluster arrive after barrier init
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_arrive_relaxed()
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
#
# Setup smem tensor A/B/SFA/SFB/C
@@ -910,18 +928,34 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
tCrB = tiled_mma.make_fragment_B(sB)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_fake = tiled_mma.make_fragment_C(
cute.append(acc_shape, self.num_acc_stage)
)
if cutlass.const_expr(self.overlapping_accum):
num_acc_stage_overlapped = 2
tCtAcc_fake = tiled_mma.make_fragment_C(
cute.append(acc_shape, num_acc_stage_overlapped)
)
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_fake = cute.make_tensor(
tCtAcc_fake.iterator,
cute.make_layout(
tCtAcc_fake.shape,
stride = (
tCtAcc_fake.stride[0],
tCtAcc_fake.stride[1],
tCtAcc_fake.stride[2],
(256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1]
)
)
)
else:
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_fake = tiled_mma.make_fragment_C(
cute.append(acc_shape, self.num_acc_stage)
)
#
# Cluster wait before tensor memory alloc
#
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
else:
self.cta_sync_barrier.arrive_and_wait()
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
#
# Specialized TMA load warp
@@ -1057,7 +1091,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
# Make SFA tmem tensor
sfa_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base),
acc_tmem_ptr + self.num_accumulator_tmem_cols,
dtype=self.sf_dtype,
)
# (MMA, MMA_M, MMA_K)
@@ -1071,9 +1105,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
# Make SFB tmem tensor
sfb_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc_base)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA),
acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols,
dtype=self.sf_dtype,
)
# (MMA, MMA_N, MMA_K)
@@ -1122,9 +1154,15 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
cur_tile_coord[2],
)
# Get accumulator stage index
if cutlass.const_expr(self.overlapping_accum):
acc_stage_index = acc_producer_state.phase ^ 1
else:
acc_stage_index = acc_producer_state.index
# Set tensor memory buffer for current tile
# (MMA, MMA_M, MMA_N)
tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)]
tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)]
# Peek (try_wait) AB buffer full for k_tile = 0
ab_consumer_state.reset_count()
@@ -1146,8 +1184,8 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0)
shifted_ptr = cute.recast_ptr(
acc_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc_base)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+ self.num_accumulator_tmem_cols
+ self.num_sfa_tmem_cols
+ offset,
dtype=self.sf_dtype,
)
@@ -1156,9 +1194,9 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
# Move in increments of 64 columns of SFB
offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
shifted_ptr = cute.recast_ptr(
acc_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc_base)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA)
acc_tmem_ptr
+ self.num_accumulator_tmem_cols
+ self.num_sfa_tmem_cols
+ offset,
dtype=self.sf_dtype,
)
@@ -1350,10 +1388,17 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
)
]
# Get accumulator stage index
if cutlass.const_expr(self.overlapping_accum):
acc_stage_index = acc_consumer_state.phase
reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False)
else:
acc_stage_index = acc_consumer_state.index
# Set tensor memory buffer for current tile
# (T2R, T2R_M, T2R_N, EPI_M, EPI_M)
tTR_tAcc = tTR_tAcc_base[
(None, None, None, None, None, acc_consumer_state.index)
(None, None, None, None, None, acc_stage_index)
]
#
@@ -1370,12 +1415,27 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
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(subtile_cnt):
real_subtile_idx = subtile_idx
if cutlass.const_expr(self.overlapping_accum):
if reverse_subtile:
real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n - 1 - subtile_idx
#
# Load accumulator from tensor memory buffer to register
#
tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)]
tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)]
cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
#
# Async arrive accumulator buffer empty ealier when overlapping_accum is enabled
#
if cutlass.const_expr(self.overlapping_accum):
if subtile_idx == self.iter_acc_early_release_in_epilogue:
# Fence for TMEM load
cute.arch.fence_view_async_tmem_load()
with cute.arch.elect_one():
acc_pipeline.consumer_release(acc_consumer_state)
acc_consumer_state.advance()
#
# Convert to C type
#
@@ -1386,7 +1446,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
#
# Store C to shared memory
#
c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage
c_buffer = (num_prev_subtiles + real_subtile_idx) % self.num_c_stage
cute.copy(
tiled_copy_r2s,
tRS_rC,
@@ -1406,7 +1466,7 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
cute.copy(
tma_atom_c,
bSG_sC[(None, c_buffer)],
bSG_gC[(None, subtile_idx)],
bSG_gC[(None, real_subtile_idx)],
)
# Fence and barrier to make sure shared memory store is visible to TMA store
c_pipeline.producer_commit()
@@ -1416,9 +1476,10 @@ class Sm100BlockScaledPersistentDenseGemmKernel:
#
# Async arrive accumulator buffer empty
#
with cute.arch.elect_one():
acc_pipeline.consumer_release(acc_consumer_state)
acc_consumer_state.advance()
if cutlass.const_expr(not self.overlapping_accum):
with cute.arch.elect_one():
acc_pipeline.consumer_release(acc_consumer_state)
acc_consumer_state.advance()
#
# Advance to next tile
@@ -2286,6 +2347,7 @@ def run(
c_tensor,
max_active_clusters,
current_stream,
options=f"--opt-level 2",
)
# Compute reference result
examples/python/CuTeDSL/blackwell/dense_gemm.py
+5 -4
View File
@@ -36,6 +36,7 @@ import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils.blackwell_helpers as sm100_utils
@@ -536,6 +537,7 @@ class DenseGemmKernel:
consumer_group=ab_pipeline_consumer_group,
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
).make_participants()
# Initialize acc_pipeline (barrier) and states
@@ -549,6 +551,7 @@ class DenseGemmKernel:
producer_group=acc_pipeline_producer_group,
consumer_group=acc_pipeline_consumer_group,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
acc_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_acc_stage
@@ -569,8 +572,7 @@ class DenseGemmKernel:
)
# Cluster arrive after barrier init
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_arrive_relaxed()
pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True)
#
# Setup smem tensor A/B/C
@@ -686,8 +688,7 @@ class DenseGemmKernel:
#
# Cluster wait before tensor memory alloc
#
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk)
# Alloc tensor memory buffer
tmem.allocate(self.num_tmem_alloc_cols)
examples/python/CuTeDSL/blackwell/dense_gemm_alpha_beta_persistent.py
+8 -12
View File
@@ -38,6 +38,7 @@ import cutlass.cute.testing as testing
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
import cutlass.utils.blackwell_helpers as sm100_utils
from cutlass.cute.nvgpu import cpasync, tcgen05
@@ -232,9 +233,8 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
)
)
# Set barrier id for cta sync, epilogue sync and tmem ptr sync
self.cta_sync_bar_id = 1
self.epilog_sync_bar_id = 2
self.tmem_alloc_sync_bar_id = 3
self.epilog_sync_bar_id = 1
self.tmem_alloc_sync_bar_id = 2
self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
def _setup_attributes(self):
@@ -635,6 +635,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
consumer_group=ab_pipeline_consumer_group,
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Initialize acc_pipeline (barrier) and states
@@ -651,6 +652,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
producer_group=acc_pipeline_producer_group,
consumer_group=acc_pipeline_consumer_group,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Load C pipeline
@@ -665,6 +667,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
producer_group=c_producer_group,
consumer_group=c_consumer_group,
tx_count=self.tma_c_load_bytes,
defer_sync=True,
)
tmem_alloc_barrier = pipeline.NamedBarrier(
@@ -681,8 +684,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
)
# Cluster arrive after barrier init
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_arrive_relaxed()
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
#
# Setup smem tensor A/B/C/D
@@ -796,9 +798,6 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
# Named barriers
#
cta_sync_barrier = pipeline.NamedBarrier(
self.cta_sync_bar_id, self.threads_per_cta
)
epilog_sync_barrier = pipeline.NamedBarrier(
self.epilog_sync_bar_id, 32 * len(self.epilog_warp_ids)
)
@@ -806,10 +805,7 @@ class SM100PersistentDenseGemmAlphaBetaKernel:
#
# Cluster wait before tensor memory alloc
#
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
else:
cta_sync_barrier.arrive_and_wait()
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
#
# Specialized TMA load warp
examples/python/CuTeDSL/blackwell/dense_gemm_persistent.py
+315 -175
View File
@@ -29,16 +29,14 @@
import argparse
from typing import Optional, Tuple, Type, Union
import torch
import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
import cutlass.cute.testing as testing
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.pipeline import pipeline_init_arrive, pipeline_init_wait
from cutlass.cute.nvgpu import cpasync, tcgen05
"""
@@ -152,10 +150,10 @@ def _compute_stages(
num_c_stage = 2 if use_tma_store else 0
# Calculate smem layout and size for one stage of A, B, and C with 1-stage
a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(
a_smem_layout_stage_one = utils.sm100.make_smem_layout_a(
tiled_mma, mma_tiler_mnk, a_dtype, 1
)
b_smem_layout_staged_one = sm100_utils.make_smem_layout_b(
b_smem_layout_staged_one = utils.sm100.make_smem_layout_b(
tiled_mma, mma_tiler_mnk, b_dtype, 1
)
@@ -229,14 +227,14 @@ class PersistentDenseGemmKernel:
- Cluster shape M must be multiple of 2 if use_2cta_instrs=True
- Cluster shape M/N must be positive and power of 2, total cluster size <= 16
Example:
>>> gemm = PersistentDenseGemmKernel(
... acc_dtype=cutlass.Float32,
... use_2cta_instrs=True,
... mma_tiler_mn=(128, 128),
... cluster_shape_mn=(2, 2)
... )
>>> gemm(a_tensor, b_tensor, c_tensor, max_active_clusters, stream)
**Example:**
gemm = PersistentDenseGemmKernel(
acc_dtype=cutlass.Float32,
use_2cta_instrs=True,
mma_tiler_mn=(128, 128),
cluster_shape_mn=(2, 2)
)
gemm(a, b, c, max_active_clusters, stream)
"""
def __init__(
@@ -316,7 +314,7 @@ class PersistentDenseGemmKernel:
- Computing tensor memory allocation columns
"""
# Configure tiled mma
tiled_mma = sm100_utils.make_trivial_tiled_mma(
tiled_mma = utils.sm100.make_trivial_tiled_mma(
self.a_dtype,
self.a_major_mode,
self.b_major_mode,
@@ -353,7 +351,7 @@ class PersistentDenseGemmKernel:
# Compute epilogue subtile
if cutlass.const_expr(self.use_tma_store):
self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
self.epi_tile = utils.sm100.compute_epilogue_tile_shape(
self.cta_tile_shape_mnk,
self.use_2cta_instrs,
self.c_layout,
@@ -364,7 +362,7 @@ class PersistentDenseGemmKernel:
c_smem_layout = None
if cutlass.const_expr(self.use_tma_store):
c_smem_layout = sm100_utils.make_smem_layout_epi(
c_smem_layout = utils.sm100.make_smem_layout_epi(
self.c_dtype, self.c_layout, self.epi_tile, 1
)
@@ -382,16 +380,16 @@ class PersistentDenseGemmKernel:
)
# Compute A/B/C shared memory layout
self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
self.a_smem_layout_staged = utils.sm100.make_smem_layout_a(
tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage
)
self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
self.b_smem_layout_staged = utils.sm100.make_smem_layout_b(
tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage
)
self.c_smem_layout_staged = None
if self.use_tma_store:
self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(
self.c_smem_layout_staged = utils.sm100.make_smem_layout_epi(
self.c_dtype, self.c_layout, self.epi_tile, self.num_c_stage
)
@@ -447,7 +445,7 @@ class PersistentDenseGemmKernel:
# Setup attributes that dependent on gemm inputs
self._setup_attributes()
tiled_mma = sm100_utils.make_trivial_tiled_mma(
tiled_mma = utils.sm100.make_trivial_tiled_mma(
self.a_dtype,
self.a_major_mode,
self.b_major_mode,
@@ -458,7 +456,7 @@ class PersistentDenseGemmKernel:
atom_thr_size = cute.size(tiled_mma.thr_id.shape)
# Setup TMA load for A
a_op = sm100_utils.cluster_shape_to_tma_atom_A(
a_op = utils.sm100.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))
@@ -475,7 +473,7 @@ class PersistentDenseGemmKernel:
)
# Setup TMA load for B
b_op = sm100_utils.cluster_shape_to_tma_atom_B(
b_op = utils.sm100.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))
@@ -614,6 +612,7 @@ class PersistentDenseGemmKernel:
consumer_group=ab_pipeline_consumer_group,
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
).make_participants()
# Initialize acc_pipeline (barrier) and states
@@ -630,6 +629,7 @@ class PersistentDenseGemmKernel:
producer_group=acc_pipeline_producer_group,
consumer_group=acc_pipeline_consumer_group,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
tmem_alloc_barrier = pipeline.NamedBarrier(
@@ -652,8 +652,7 @@ class PersistentDenseGemmKernel:
)
# Cluster arrive after barrier init
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_arrive_relaxed()
pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True)
#
# Setup smem tensor A/B/C
@@ -761,10 +760,7 @@ class PersistentDenseGemmKernel:
#
# Cluster wait before tensor memory alloc
#
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
else:
cute.arch.sync_threads()
pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk)
#
# Specialized TMA load warp
@@ -1297,7 +1293,7 @@ class PersistentDenseGemmKernel:
:rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
"""
# Make tiledCopy for tensor memory load
copy_atom_t2r = sm100_utils.get_tmem_load_op(
copy_atom_t2r = utils.sm100.get_tmem_load_op(
self.cta_tile_shape_mnk,
self.c_layout,
self.c_dtype,
@@ -1354,7 +1350,7 @@ class PersistentDenseGemmKernel:
- tRS_sC: The partitioned tensor C (smem destination)
:rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
"""
copy_atom_r2s = sm100_utils.get_smem_store_op(
copy_atom_r2s = utils.sm100.get_smem_store_op(
self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r
)
tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
@@ -1617,97 +1613,187 @@ class PersistentDenseGemmKernel:
is_valid = False
return is_valid
def can_implement(self, a: cute.Tensor, b: cute.Tensor, c: cute.Tensor) -> bool:
"""Check if the given tensors can be implemented by this kernel.
def can_implement(
self,
mnkl: Tuple[int, int, int, int],
ab_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
a_major: str,
b_major: str,
c_major: str,
) -> bool:
"""
Determine if the given tensor configuration can be implemented by this kernel.
:param a: Input tensor A
:type a: cute.Tensor
:param b: Input tensor B
:type b: cute.Tensor
:param c: Output tensor C
:type c: cute.Tensor
:return: True if the gemm supports the given config, False otherwise
:param mnkl: Problem size as a tuple (M, N, K, L).
:type mnkl: Tuple[int, int, int, int]
:param ab_dtype: Data type for input tensors A and B.
:type ab_dtype: Type[cutlass.Numeric]
:param c_dtype: Data type for output tensor C.
:type c_dtype: Type[cutlass.Numeric]
:param a_major: Major dimension of the A tensor layout ("m" or "k").
:type a_major: str
:param b_major: Major dimension of the B tensor layout ("n" or "k").
:type b_major: str
:param c_major: Major dimension of the C tensor layout ("m" or "n").
:type c_major: str
:return: True if the kernel supports the given configuration, False otherwise.
:rtype: bool
"""
m, n, k, l = a.shape[0], b.shape[0], a.shape[1], a.shape[2]
# infer a_major, b_major, c_major
is_m_major_a = utils.LayoutEnum.from_tensor(a).is_m_major_a()
is_n_major_b = utils.LayoutEnum.from_tensor(b).is_n_major_b()
is_m_major_c = utils.LayoutEnum.from_tensor(c).is_m_major_c()
a_major = "m" if is_m_major_a else "k"
b_major = "n" if is_n_major_b else "k"
c_major = "m" if is_m_major_c else "n"
can_implement = True
# Skip unsupported types
if not self.is_valid_dtypes(a.element_type, c.element_type):
can_implement = False
if not self.is_valid_dtypes(ab_dtype, c_dtype):
return False
# Skip invalid mma tile shape and cluster shape
if not self.is_valid_mma_tiler_and_cluster_shape():
can_implement = False
return False
# Unpack mnkl for clarity in calling the epilog check
m, n, k, l = mnkl
# Skip illegal problem shape for load/store alignment
if not self.is_valid_tensor_alignment(
m, n, k, l, a.element_type, c.element_type, a_major, b_major, c_major
m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major
):
can_implement = False
return False
# Skip invalid epilogue store option
if not self.is_valid_epilog_store_option(m, n):
can_implement = False
return False
return can_implement
return True
def create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype):
torch.manual_seed(1111)
@cute.jit
def bmm(
gemm_op: cutlass.Constexpr,
a: cute.Tensor, # (l, m, k)
b: cute.Tensor, # (l, k, n)
c: cute.Tensor, # (l, m, n)
max_active_clusters: cutlass.Constexpr,
stream: cuda.CUstream,
epilogue_op: cutlass.Constexpr = lambda x: x,
):
"""
Wrapper API for persistent GEMM kernel to follow the convention of PyTorch's batch matrix-multiply (bmm).
a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype)
b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype)
c_torch_cpu = cutlass_torch.matrix(l, m, n, c_major == "m", c_dtype)
Internally, the tensors are permuted to match CuTe's convention:
- a: (m, k, l)
- b: (n, k, l)
- c: (m, n, l)
a_tensor, _ = cutlass_torch.cute_tensor_like(
a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16
:param gemm_op: Kernel operation, expects (a, b, c, max_active_clusters, stream, epilogue_op)
:type gemm_op: cutlass.Constexpr
:param a: Input tensor of shape (l, m, k)
:type a: cute.Tensor
:param b: Input tensor of shape (l, k, n)
:type b: cute.Tensor
:param c: Output tensor of shape (l, m, n)
:type c: cute.Tensor
:param max_active_clusters: Maximum number of hardware clusters to launch
:type max_active_clusters: cutlass.Constexpr
:param epilogue_op: Optional elementwise lambda function to apply per output element, defaults to identity
:type epilogue_op: cutlass.Constexpr, optional
"""
# (l,m,k) -> (m,k,l)
a = cute.make_tensor(a.iterator, cute.select(a.layout, mode=[1, 2, 0]))
# (l,k,n) -> (n,k,l)
b = cute.make_tensor(b.iterator, cute.select(b.layout, mode=[2, 1, 0]))
# (l,m,n) -> (m,n,l)
c = cute.make_tensor(c.iterator, cute.select(c.layout, mode=[1, 2, 0]))
gemm_op(a, b, c, max_active_clusters, stream, epilogue_op)
def compile_bmm(
gemm_op: PersistentDenseGemmKernel,
a_dtype: Type[cutlass.Numeric],
b_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
a_major: str,
b_major: str,
c_major: str,
max_active_clusters: cutlass.Constexpr,
stream: cuda.CUstream,
epilogue_op: cutlass.Constexpr = lambda x: x,
options: str = "",
):
from cutlass.cute.runtime import make_fake_compact_tensor
a_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
b_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
c_shape = (cute.sym_int(), cute.sym_int(divisibility=16), cute.sym_int())
if a_major == "k":
a_order = (2, 1, 0) # k is leading dimension
elif a_major == "m":
a_order = (2, 0, 1) # m is leading dimension
if b_major == "n":
b_order = (2, 1, 0) # n is leading dimension
elif b_major == "k":
b_order = (2, 0, 1) # k is leading dimension
if c_major == "n":
c_order = (2, 1, 0) # n is leading dimension
elif c_major == "m":
c_order = (2, 0, 1) # m is leading dimension
a = make_fake_compact_tensor(
a_dtype, a_shape, stride_order=a_order, assumed_align=16
)
b_tensor, _ = cutlass_torch.cute_tensor_like(
b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16
b = make_fake_compact_tensor(
b_dtype, b_shape, stride_order=b_order, assumed_align=16
)
c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like(
c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16
c = make_fake_compact_tensor(
c_dtype, c_shape, stride_order=c_order, assumed_align=16
)
return cute.compile(
bmm, gemm_op, a, b, c, max_active_clusters, stream, epilogue_op, options=options
)
def prepare_tensors(
mnkl: Tuple[int, int, int, int],
ab_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
a_major: str,
b_major: str,
c_major: str,
init_random: bool = True,
):
import torch
from cutlass.torch import dtype as torch_dtype
m, n, k, l = mnkl
if a_major == "k":
a = torch.empty((l, m, k), dtype=torch.float32, device="cuda")
elif a_major == "m":
a = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute(0, 2, 1)
if b_major == "n":
b = torch.empty((l, k, n), dtype=torch.float32, device="cuda")
elif b_major == "k":
b = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute(0, 2, 1)
if c_major == "n":
c = torch.empty((l, m, n), dtype=torch.float32, device="cuda")
elif c_major == "m":
c = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute(0, 2, 1)
if init_random:
a.random_(-2, 3)
b.random_(-2, 3)
c.random_(-2, 3)
return (
a_tensor,
b_tensor,
c_tensor,
a_torch_cpu,
b_torch_cpu,
c_torch_cpu,
c_torch_gpu,
a.to(dtype=torch_dtype(ab_dtype)),
b.to(dtype=torch_dtype(ab_dtype)),
c.to(dtype=torch_dtype(c_dtype)),
)
def compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance):
# Copy gpu result back
kernel_result = c_torch_gpu.cpu()
# Compute reference result
ref = torch.einsum(
"mkl,nkl->mnl",
a_torch_cpu.to(dtype=torch.float32),
b_torch_cpu.to(dtype=torch.float32),
)
# Convert ref to c_dtype
_, ref_torch_gpu = cutlass_torch.cute_tensor_like(
ref, c_dtype, is_dynamic_layout=True, assumed_align=16
)
ref_result = ref_torch_gpu.cpu()
# Assert close results
torch.testing.assert_close(kernel_result, ref_result, atol=tolerance, rtol=1e-05)
def run(
mnkl: Tuple[int, int, int, int],
ab_dtype: Type[cutlass.Numeric],
@@ -1725,48 +1811,55 @@ def run(
iterations: int = 1,
skip_ref_check: bool = False,
use_cold_l2: bool = False,
use_tvm_ffi: bool = False,
benchmark: bool = False,
**kwargs,
):
"""Execute a persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking.
"""
Execute a persistent batched dense GEMM operation on Blackwell architecture with performance benchmarking.
This function prepares input tensors, configures and launches the persistent GEMM kernel,
optionally performs reference validation, and benchmarks the execution performance.
Prepares input tensors, configures and launches the persistent GEMM kernel,
optionally performs reference validation, and benchmarks execution.
:param mnkl: Problem size (M, N, K, L)
:param mnkl: Problem size as a tuple (M, N, K, L).
:type mnkl: Tuple[int, int, int, int]
:param ab_dtype: Data type for input tensors A and B
:param ab_dtype: Data type for input tensors A and B.
:type ab_dtype: Type[cutlass.Numeric]
:param c_dtype: Data type for output tensor C
:param c_dtype: Data type for output tensor C.
:type c_dtype: Type[cutlass.Numeric]
:param acc_dtype: Data type for accumulation during matrix multiplication
:param acc_dtype: Accumulator data type for the matrix multiplication.
:type acc_dtype: Type[cutlass.Numeric]
:param a_major/b_major/c_major: Memory layout of tensor A/B/C
:type a_major/b_major/c_major: str
:param mma_tiler_mn: MMA tiling size. If not specified in the decorator parameters, the autotuner will use the
default value of (256, 256). Otherwise, the autotuner will use the value specified in the decorator parameters.
:param a_major: Memory layout of tensor A.
:type a_major: str
:param b_major: Memory layout of tensor B.
:type b_major: str
:param c_major: Memory layout of tensor C.
:type c_major: str
:param mma_tiler_mn: MMA tiling size (M, N), defaults to (256, 256).
:type mma_tiler_mn: Tuple[int, int], optional
:param cluster_shape_mn: Cluster shape. If not specified in the decorator parameters, the autotuner will use the
default value of (2, 1). Otherwise, the autotuner will use the value specified in the decorator parameters.
:param cluster_shape_mn: Cluster shape (M, N), defaults to (2, 1).
:type cluster_shape_mn: Tuple[int, int], optional
:param use_2cta_instrs: Whether to use 2CTA instructions. If not specified in the decorator parameters, the autotuner
will use the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters.
:param use_2cta_instrs: Whether to use 2CTA MMA instructions, defaults to True.
:type use_2cta_instrs: bool, optional
:param use_tma_store: Whether to use TMA store. If not specified in the decorator parameters, the autotuner will use
the default value of True. Otherwise, the autotuner will use the value specified in the decorator parameters.
:param use_tma_store: Whether to use TMA store, defaults to True.
:type use_tma_store: bool, optional
:param tolerance: Tolerance value for reference validation comparison, defaults to 1e-01
:param tolerance: Tolerance for reference validation, defaults to 1e-01.
:type tolerance: float, optional
:param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0
:param warmup_iterations: Number of warmup iterations before benchmarking, defaults to 0.
:type warmup_iterations: int, optional
:param iterations: Number of benchmark iterations to run, defaults to 1
:param iterations: Number of benchmark iterations to run, defaults to 1.
:type iterations: int, optional
:param skip_ref_check: Whether to skip reference result validation, defaults to False
:param skip_ref_check: Whether to skip reference result validation, defaults to False.
:type skip_ref_check: bool, optional
:param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False
:param use_cold_l2: Whether to use circular buffer strategy to ensure cold L2 cache, defaults to False.
:type use_cold_l2: bool, optional
:raises RuntimeError: If CUDA GPU is not available
:raises ValueError: If the configuration is invalid or unsupported by the kernel
:return: Execution time of the GEMM kernel
:param use_tvm_ffi: Whether to use TVM FFI for the kernel, defaults to False.
:type use_tvm_ffi: bool, optional
:param benchmark: Whether to only benchmark the kernel, defaults to False.
:type benchmark: bool, optional
:raises RuntimeError: If CUDA GPU is not available.
:raises ValueError: If the configuration is invalid or unsupported by the kernel.
:return: Execution time of the GEMM kernel.
:rtype: float
"""
print("Running Blackwell Persistent Dense GEMM test with:")
@@ -1781,9 +1874,24 @@ def run(
print(f"Iterations: {iterations}")
print(f"Skip reference checking: {skip_ref_check}")
print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}")
print(f"Use TVM FFI: {'True' if use_tvm_ffi else 'False'}")
# Unpack parameters
m, n, k, l = mnkl
import torch
from cutlass.torch import dtype as torch_dtype
# Build GEMM object
gemm = PersistentDenseGemmKernel(
acc_dtype, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, use_tma_store
)
can_implement = gemm.can_implement(
mnkl, ab_dtype, c_dtype, a_major, b_major, c_major
)
if not can_implement:
raise testing.CantImplementError(
f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, "
f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, "
f"use_tma_store = {use_tma_store}"
)
if not torch.cuda.is_available():
raise RuntimeError("GPU is required to run this example!")
@@ -1793,59 +1901,75 @@ def run(
# Get the raw stream pointer as a CUstream
current_stream = cuda.CUstream(torch_stream.cuda_stream)
a_tensor, b_tensor, c_tensor, a_torch_cpu, b_torch_cpu, c_torch_cpu, c_torch_gpu = (
create_tensors(l, m, n, k, a_major, b_major, c_major, ab_dtype, c_dtype)
)
# Build GEMM object
gemm = PersistentDenseGemmKernel(
acc_dtype, use_2cta_instrs, mma_tiler_mn, cluster_shape_mn, use_tma_store
)
# Check if configuration can be implemented
can_implement = gemm.can_implement(a_tensor, b_tensor, c_tensor)
if not can_implement:
raise ValueError(
f"The current config which is invalid/unsupported: use_2cta_instrs = {use_2cta_instrs}, "
f"mma_tiler_mn = {mma_tiler_mn}, cluster_shape_mn = {cluster_shape_mn}, "
f"use_tma_store = {use_tma_store}"
)
max_active_clusters = utils.HardwareInfo().get_max_active_clusters(
cluster_shape_mn[0] * cluster_shape_mn[1]
)
compiled_gemm = cute.compile(
gemm, a_tensor, b_tensor, c_tensor, max_active_clusters, current_stream
options = []
if use_tvm_ffi:
options.append("--enable-tvm-ffi")
compiled_fn = compile_bmm(
gemm,
ab_dtype,
ab_dtype,
c_dtype,
a_major,
b_major,
c_major,
max_active_clusters,
current_stream,
options=",".join(options),
)
# Run and verify BMM with torch
a, b, c = prepare_tensors(mnkl, ab_dtype, c_dtype, a_major, b_major, c_major)
if not skip_ref_check:
compiled_gemm(a_tensor, b_tensor, c_tensor, current_stream)
compare(a_torch_cpu, b_torch_cpu, c_torch_gpu, c_dtype, tolerance)
# Use small random number for deterministic result for reference check
compiled_fn(a, b, c, torch_stream)
# Manually quantize to be comparable
ref = (
torch.bmm(a.to(dtype=torch.float32), b.to(dtype=torch.float32))
.to(dtype=torch_dtype(c_dtype))
.to(dtype=torch.float32)
)
torch.testing.assert_close(
c.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03
)
if not benchmark:
return 0
def generate_tensors():
a_tensor, _ = cutlass_torch.cute_tensor_like(
a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16
init_normal = ab_dtype not in [cutlass.Int8, cutlass.Uint8]
a, b, c = prepare_tensors(
mnkl,
ab_dtype,
c_dtype,
a_major,
b_major,
c_major,
init_random=not init_normal,
)
b_tensor, _ = cutlass_torch.cute_tensor_like(
b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16
)
c_tensor, _ = cutlass_torch.cute_tensor_like(
c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16
)
return testing.JitArguments(a_tensor, b_tensor, c_tensor, current_stream)
return testing.JitArguments(a, b, c, torch_stream)
workspace_count = 1
if use_cold_l2:
one_workspace_bytes = (
a_torch_cpu.numel() * a_torch_cpu.element_size()
+ b_torch_cpu.numel() * b_torch_cpu.element_size()
+ c_torch_cpu.numel() * c_torch_cpu.element_size()
a.numel() * a.element_size()
+ b.numel() * b.element_size()
+ c.numel() * c.element_size()
)
workspace_count = testing.get_workspace_count(
one_workspace_bytes, warmup_iterations, iterations
)
exec_time = testing.benchmark(
compiled_gemm,
# Return execution time in microseconds
return testing.benchmark(
compiled_fn,
workspace_generator=generate_tensors,
workspace_count=workspace_count,
stream=current_stream,
@@ -1853,18 +1977,17 @@ def run(
iterations=iterations,
)
return exec_time # Return execution time in microseconds
def _parse_comma_separated_ints(s: str) -> Tuple[int, ...]:
try:
return tuple(int(x.strip()) for x in s.split(","))
except ValueError:
raise argparse.ArgumentTypeError(
"Invalid format. Expected comma-separated integers."
)
if __name__ == "__main__":
def parse_comma_separated_ints(s: str) -> Tuple[int, ...]:
try:
return tuple(int(x.strip()) for x in s.split(","))
except ValueError:
raise argparse.ArgumentTypeError(
"Invalid format. Expected comma-separated integers."
)
def prepare_parser():
parser = argparse.ArgumentParser(
description="Example of Dense Persistent GEMM on Blackwell."
@@ -1872,19 +1995,13 @@ if __name__ == "__main__":
parser.add_argument(
"--mnkl",
type=parse_comma_separated_ints,
type=_parse_comma_separated_ints,
default=(256, 256, 512, 1),
help="mnkl dimensions (comma-separated)",
)
parser.add_argument(
"--mma_tiler_mn",
type=parse_comma_separated_ints,
default=(128, 128),
help="Mma tile shape (comma-separated)",
)
parser.add_argument(
"--cluster_shape_mn",
type=parse_comma_separated_ints,
type=_parse_comma_separated_ints,
default=(1, 1),
help="Cluster shape (comma-separated)",
)
@@ -1905,6 +2022,9 @@ if __name__ == "__main__":
parser.add_argument(
"--tolerance", type=float, default=1e-01, help="Tolerance for validation"
)
parser.add_argument(
"--benchmark", action="store_true", help="Only benchmark the kernel"
)
parser.add_argument(
"--warmup_iterations", type=int, default=0, help="Warmup iterations"
)
@@ -1923,6 +2043,24 @@ if __name__ == "__main__":
default=False,
help="Use circular buffer tensor sets to ensure L2 cold cache",
)
parser.add_argument(
"--use_tvm_ffi",
action="store_true",
default=False,
help="Enable TVM FFI for the kernel, defaults to False using CuTe DSL's native runtime",
)
return parser
if __name__ == "__main__":
parser = prepare_parser()
parser.add_argument(
"--mma_tiler_mn",
type=_parse_comma_separated_ints,
default=(128, 128),
help="Mma tile shape (comma-separated)",
)
args = parser.parse_args()
@@ -1952,5 +2090,7 @@ if __name__ == "__main__":
args.iterations,
args.skip_ref_check,
args.use_cold_l2,
args.use_tvm_ffi,
args.benchmark,
)
print("PASS")
examples/python/CuTeDSL/blackwell/dense_gemm_software_pipeline.py
+5 -4
View File
@@ -36,6 +36,7 @@ import cutlass
import cutlass.cute as cute
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils.blackwell_helpers as sm100_utils
@@ -535,6 +536,7 @@ class DenseGemmKernel:
consumer_group=ab_pipeline_consumer_group,
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
).make_participants()
# Initialize acc_pipeline (barrier) and states
@@ -549,6 +551,7 @@ class DenseGemmKernel:
producer_group=acc_pipeline_producer_group,
consumer_group=acc_pipeline_consumer_group,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
acc_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_acc_stage
@@ -569,8 +572,7 @@ class DenseGemmKernel:
)
# Cluster arrive after barrier init
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_arrive_relaxed()
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
#
# Setup smem tensor A/B/C
@@ -686,8 +688,7 @@ class DenseGemmKernel:
#
# Cluster wait before tensor memory alloc
#
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
# Alloc tensor memory buffer
tmem.allocate(self.num_tmem_alloc_cols)
examples/python/CuTeDSL/blackwell/fmha.py
+5 -3
View File
@@ -48,9 +48,11 @@ import cutlass.cute.testing as testing
from cutlass.cute.runtime import from_dlpack
from cutlass.cute.typing import Int32, Int64, Float32
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(current_dir, ".."))
from utils import fmha_helpers as fmha_utils
if __name__ == "__main__":
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(current_dir, ".."))
from helpers import fmha_helpers as fmha_utils
"""
A fused multi-head attention (FMHA) example for the NVIDIA Blackwell SM100 architecture using CUTE DSL
examples/python/CuTeDSL/blackwell/fmha_bwd.py
+5 -3
View File
@@ -49,9 +49,11 @@ import cutlass.utils.blackwell_helpers as sm100_utils
from cutlass.cute.runtime import from_dlpack
from cutlass.cute.typing import Int32, Float32, Float8E4M3FN, Float16, BFloat16, Boolean
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(current_dir, ".."))
from utils import fmha_helpers as fmha_utils
if __name__ == "__main__":
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(current_dir, ".."))
from helpers import fmha_helpers as fmha_utils
"""
A fused multi-head attention (FMHA) backward pass example for the NVIDIA Blackwell SM100 architecture using CUTE DSL
examples/python/CuTeDSL/blackwell/grouped_blockscaled_gemm.py
+9 -15
View File
@@ -40,6 +40,7 @@ from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.cute.runtime import from_dlpack
@@ -177,22 +178,18 @@ class Sm100GroupedBlockScaledGemmKernel:
self.threads_per_cta = 32 * len(
(self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
)
# Set barrier for cta sync, epilogue sync and tmem ptr sync
self.cta_sync_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=self.threads_per_cta,
)
# Set barrier for epilogue sync and tmem ptr sync
self.epilog_sync_barrier = pipeline.NamedBarrier(
barrier_id=2,
barrier_id=1,
num_threads=32 * len(self.epilog_warp_id),
)
self.tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=3,
barrier_id=2,
num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
)
# Barrier used by MMA/TMA warps to signal A/B tensormap initialization completion
self.tensormap_ab_init_barrier = pipeline.NamedBarrier(
barrier_id=4,
barrier_id=3,
num_threads=64,
)
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
@@ -646,6 +643,7 @@ class Sm100GroupedBlockScaledGemmKernel:
cluster=(*self.cluster_shape_mn, 1),
smem=self.shared_storage.size_in_bytes(),
stream=stream,
min_blocks_per_mp=1,
)
return
@@ -781,11 +779,9 @@ class Sm100GroupedBlockScaledGemmKernel:
cute.arch.mbarrier_init(
tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads
)
cute.arch.mbarrier_init_fence()
# Cluster arrive after barrier init
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_arrive_relaxed()
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
#
# Setup smem tensor A/B/SFA/SFB/C
@@ -944,10 +940,7 @@ class Sm100GroupedBlockScaledGemmKernel:
#
# Cluster wait before tensor memory alloc
#
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
else:
self.cta_sync_barrier.arrive_and_wait()
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
#
# Get tensormap buffer address
@@ -2894,6 +2887,7 @@ def run(
tensor_of_tensormap,
max_active_clusters,
current_stream,
options=f"--opt-level 2",
)
# reference check
examples/python/CuTeDSL/blackwell/grouped_gemm.py
+7 -15
View File
@@ -39,6 +39,7 @@ import cutlass.cute as cute
import cutlass.cute.testing as testing
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.torch as cutlass_torch
@@ -153,22 +154,18 @@ class GroupedGemmKernel:
self.threads_per_cta = 32 * len(
(self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
)
# Set barrier for cta sync, epilog sync, tmem ptr sync and tensormap update sync
self.cta_sync_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=self.threads_per_cta,
)
# Set barrier for epilog sync, tmem ptr sync and tensormap update sync
self.epilog_sync_barrier = pipeline.NamedBarrier(
barrier_id=2,
barrier_id=1,
num_threads=32 * len(self.epilog_warp_id),
)
self.tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=3,
barrier_id=2,
num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
)
# Barrier used by MMA/TMA warps to signal A/B tensormap initialization completion
self.tensormap_ab_init_barrier = pipeline.NamedBarrier(
barrier_id=4,
barrier_id=3,
num_threads=32 * (len(self.epilog_warp_id) + 1),
)
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
@@ -586,11 +583,9 @@ class GroupedGemmKernel:
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
)
cute.arch.mbarrier_init_fence()
# Cluster arrive after barrier init
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_arrive_relaxed()
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
#
# Setup smem tensor A/B/C
@@ -718,10 +713,7 @@ class GroupedGemmKernel:
#
# Cluster wait before tensor memory alloc
#
if cute.size(self.cluster_shape_mn) > 1:
cute.arch.cluster_wait()
else:
self.cta_sync_barrier.arrive_and_wait()
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
#
# Get tensormap buffer address
examples/python/CuTeDSL/blackwell/grouped_mixed_input_gemm.py
+3202
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/mamba2_ssd/mamba2_ssd.py
+23 -9
View File
@@ -27,6 +27,8 @@
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import os
import sys
import argparse
from typing import List, Type, Tuple, Optional
import cuda.bindings.driver as cuda
@@ -39,21 +41,22 @@ import cutlass.cute as cute
import cutlass.cute.testing as testing
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils.blackwell_helpers as sm100_utils
from cutlass.cute.runtime import from_dlpack
import sys
from pathlib import Path
if __name__ == "__main__":
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(current_dir, "../.."))
sys.path.append(str(Path(__file__).resolve().parent))
from mamba2_ssd_reference import (
from blackwell.mamba2_ssd.mamba2_ssd_reference import (
ssd_reference_fp32_all,
ssd_reference_lowprecision_intermediates,
analyze_relative_diffs,
)
from mamba2_ssd_tile_scheduler import (
from blackwell.mamba2_ssd.mamba2_ssd_tile_scheduler import (
Mamba2SSDTileSchedulerParams,
Mamba2SSDTileScheduler,
)
@@ -811,12 +814,10 @@ class SSDKernel:
)
# Cluster arrive after barrier init
if cute.size(self.cluster_shape_mnk) > 1:
cute.arch.cluster_arrive_relaxed()
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True)
# Cluster wait before tmem alloc
if cute.size(self.cluster_shape_mnk) > 1:
cute.arch.cluster_wait()
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk)
tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=0,
@@ -2574,6 +2575,7 @@ class SSDKernel:
consumer_group=x_consumer_group,
tx_count=self.num_x_load_bytes,
barrier_storage=x_full_mbar_ptr,
defer_sync=True,
)
else:
x_consumer_group_umma = pipeline.CooperativeGroup(
@@ -2590,6 +2592,7 @@ class SSDKernel:
consumer_group_async=x_consumer_group_async,
tx_count=self.num_x_load_bytes,
barrier_storage=x_full_mbar_ptr,
defer_sync=True,
)
def make_and_init_b_pipeline(self, b_full_mbar_ptr):
@@ -2609,6 +2612,7 @@ class SSDKernel:
consumer_group_async=b_consumer_group_async,
tx_count=self.num_b_load_bytes,
barrier_storage=b_full_mbar_ptr,
defer_sync=True,
)
def make_and_init_c_pipeline(self, c_full_mbar_ptr):
@@ -2624,6 +2628,7 @@ class SSDKernel:
consumer_group=c_consumer_group,
tx_count=self.num_c_load_bytes,
barrier_storage=c_full_mbar_ptr,
defer_sync=True,
)
def make_and_init_deltas_pipeline(self, deltas_full_mbar_ptr):
@@ -2643,6 +2648,7 @@ class SSDKernel:
consumer_group=deltas_consumer_group,
tx_count=self.num_delta_load_bytes + self.num_cumsum_delta_load_bytes,
barrier_storage=deltas_full_mbar_ptr,
defer_sync=True,
)
def make_and_init_d_pipeline(self, d_full_mbar_ptr):
@@ -2662,6 +2668,7 @@ class SSDKernel:
consumer_group=d_consumer_group,
tx_count=self.num_d_load_bytes,
barrier_storage=d_full_mbar_ptr,
defer_sync=True,
)
def make_and_init_intra1_acc_pipeline(self, intra1_acc_full_mbar_ptr):
@@ -2676,6 +2683,7 @@ class SSDKernel:
producer_group=intra1_acc_producer_group,
consumer_group=intra1_acc_consumer_group,
barrier_storage=intra1_acc_full_mbar_ptr,
defer_sync=True,
)
def make_and_init_intra2_q_pipeline(self, intra2_q_full_mbar_ptr):
@@ -2690,6 +2698,7 @@ class SSDKernel:
producer_group=intra2_q_producer_group,
consumer_group=intra2_q_consumer_group,
barrier_storage=intra2_q_full_mbar_ptr,
defer_sync=True,
)
def make_and_init_intra2_acc_pipeline(self, intra2_acc_full_mbar_ptr):
@@ -2704,6 +2713,7 @@ class SSDKernel:
producer_group=intra2_acc_producer_group,
consumer_group=intra2_acc_consumer_group,
barrier_storage=intra2_acc_full_mbar_ptr,
defer_sync=True,
)
def make_and_init_inter1_b_pipeline(self, inter1_b_full_mbar_ptr):
@@ -2718,6 +2728,7 @@ class SSDKernel:
producer_group=inter1_b_producer_group,
consumer_group=inter1_b_consumer_group,
barrier_storage=inter1_b_full_mbar_ptr,
defer_sync=True,
)
def make_and_init_inter1_acc_pipeline(self, inter1_acc_full_mbar_ptr):
@@ -2732,6 +2743,7 @@ class SSDKernel:
producer_group=inter1_acc_producer_group,
consumer_group=inter1_acc_consumer_group,
barrier_storage=inter1_acc_full_mbar_ptr,
defer_sync=True,
)
def make_and_init_inter2_p_pipeline(self, inter2_p_full_mbar_ptr):
@@ -2746,6 +2758,7 @@ class SSDKernel:
producer_group=inter2_p_producer_group,
consumer_group=inter2_p_consumer_group,
barrier_storage=inter2_p_full_mbar_ptr,
defer_sync=True,
)
def make_and_init_inter2_acc_pipeline(self, inter2_acc_full_mbar_ptr):
@@ -2760,6 +2773,7 @@ class SSDKernel:
producer_group=inter2_acc_producer_group,
consumer_group=inter2_acc_consumer_group,
barrier_storage=inter2_acc_full_mbar_ptr,
defer_sync=True,
)
def tma_partition_for_mma_b_operand(
examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_prefill.py
+2353
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/blackwell/mixed_input_gemm.py
+74 -519
View File
@@ -27,7 +27,6 @@
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import argparse
from enum import Enum, auto
from math import log2, ceil
from typing import Optional, Union
@@ -37,9 +36,12 @@ import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.mixed_input_helpers as mixed_input_utils
from cutlass.utils.mixed_input_helpers import TransformMode
import cutlass.cute.testing as testing
from cutlass.cute.nvgpu import cpasync, tcgen05
from cutlass.cute.runtime import from_dlpack
@@ -130,15 +132,6 @@ Besides the requirements from the Blackwell dense GEMM example, there are some c
"""
class TransformMode(Enum):
"""
An enumeration for the possible transform modes of a mixed-input GEMM.
"""
ConvertOnly = auto()
ConvertScale = auto()
class MixedInputGemmKernel:
"""
Mixed-input GEMM kernel for NVIDIA Blackwell SM100 architecture.
@@ -226,7 +219,7 @@ class MixedInputGemmKernel:
+ 1
)
# Set barrier id for cta sync, epilogue sync, tmem ptr sync, and transform sync
# Set barrier id for epilogue sync, tmem ptr sync, and transform sync
self.epilog_sync_barrier = pipeline.NamedBarrier(
1, 32 * len(self.epilog_warp_id)
)
@@ -234,7 +227,6 @@ class MixedInputGemmKernel:
self.transform_sync_barrier = pipeline.NamedBarrier(
3, 32 * len(self.transform_warp_id)
)
self.cta_sync_barrier = pipeline.NamedBarrier(4, self.threads_per_cta)
self.smem_buffer_align_bytes = 1024
@@ -255,7 +247,9 @@ class MixedInputGemmKernel:
- Computing tensor memory allocation columns
"""
# Deduce where the transformed A tensor is stored, shared memory(SMEM) or tensor memory(TMEM)
self.transform_a_source = self._get_transform_a_source(self.a_major_mode)
self.transform_a_source = mixed_input_utils.get_transform_a_source(
self.a_major_mode
)
tiled_mma = sm100_utils.make_trivial_tiled_mma(
self.mma_dtype,
self.a_major_mode,
@@ -346,7 +340,7 @@ class MixedInputGemmKernel:
self.smem_layout_a,
self.smem_layout_a_transform,
self.smem_layout_b,
) = self._compute_smem_layout(
) = mixed_input_utils.compute_smem_layout(
tiled_mma,
self.mma_tiler,
self.a_dtype,
@@ -358,11 +352,20 @@ class MixedInputGemmKernel:
self.smem_layout_scale_per_stage = None
self.smem_layout_scale = None
if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale):
# Get smem layout for scale tensor
# Get scale tile shape and smem layout for scale tensor
(
self.scale_tile_shape,
self.smem_layout_scale_per_stage,
self.smem_layout_scale,
) = self.get_smem_layout_scale()
) = mixed_input_utils.get_smem_layout_scale(
self.mma_tiler,
self.use_2cta_instrs,
self.scale_granularity_m,
self.scale_granularity_k,
self.scale_major_mode,
self.a_scale_dtype,
self.num_scale_load2trans_stage,
)
def _validate_inputs(
self,
@@ -448,7 +451,12 @@ class MixedInputGemmKernel:
self.c_layout = utils.LayoutEnum.from_tensor(c)
if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale):
# Get gmem layout for scale tensor
self.gmem_layout_scale = self.get_gmem_layout_scale(a.shape)
self.gmem_layout_scale = mixed_input_utils.get_gmem_layout_scale(
a.shape,
self.scale_granularity_m,
self.scale_granularity_k,
self.scale_major_mode,
)
# Validate inputs
self._validate_inputs(a, a_scale, b, c)
@@ -466,8 +474,12 @@ class MixedInputGemmKernel:
self.transform_a_source,
)
# Set up gmem copy atoms for A, scale, and B
a_op = self._get_tma_atom_kind(self.is_a_mcast, self.use_2cta_instrs, False)
b_op = self._get_tma_atom_kind(self.is_b_mcast, self.use_2cta_instrs, True)
a_op = mixed_input_utils.get_tma_atom_kind(
self.is_a_mcast, self.use_2cta_instrs, False
)
b_op = mixed_input_utils.get_tma_atom_kind(
self.is_b_mcast, self.use_2cta_instrs, True
)
a_scale_op = a_op
# Deduce TMA copy atom and TMA tensor for A, scale, and B
smem_layout_a_per_stage = cute.slice_(self.smem_layout_a, (None, None, None, 0))
@@ -650,7 +662,6 @@ class MixedInputGemmKernel:
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,
)
@@ -730,6 +741,7 @@ class MixedInputGemmKernel:
cta_layout_vmnk=cluster_layout_vmnk,
tidx=transform_thread_idx,
mcast_mode_mn=(1, 0), # multicast for A will only happen on the M-mode
defer_sync=True,
)
# Initialize scale_load2trans pipeline, which tracks the dependencies between TMA's loading
# of scale, and the transformation of A
@@ -753,6 +765,7 @@ class MixedInputGemmKernel:
1,
0,
), # multicast for scale_a will only happen on the M-mode
defer_sync=True,
)
# Initialize transform2mma pipeline, which tracks the dependencies between the transformation
# of A and MMA's consumption of transformed A
@@ -766,6 +779,7 @@ class MixedInputGemmKernel:
),
consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Initialize pipeline for tensor B load to MMA
# MMA warp informs TMA warp to proceed to load next tile of B tensor
@@ -779,6 +793,7 @@ class MixedInputGemmKernel:
tx_count=self.num_tma_load_bytes_b,
cta_layout_vmnk=cluster_layout_vmnk,
mcast_mode_mn=(0, 1), # multicast for B will only happen on the N-mode
defer_sync=True,
)
# Initialize accumulator pipeline, which tracks the dependencies between
# MMA's computation of accumulators and epilogue warps' consumption of accumulators
@@ -790,6 +805,7 @@ class MixedInputGemmKernel:
pipeline.Agent.Thread, cta_v_size * len(self.epilog_warp_id)
),
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Tensor memory dealloc barrier init
@@ -802,8 +818,7 @@ class MixedInputGemmKernel:
)
# Cluster arrive after barrier init
if cutlass.const_expr(cute.size(self.cluster_shape_mn) > 1):
cute.arch.cluster_arrive_relaxed()
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
# Setup smem tensor A/scale/B/C
sC = (
@@ -897,7 +912,7 @@ class MixedInputGemmKernel:
cute.dice(self.mma_tiler, (1, None, 1))
)
# Setup copy atom to store transformed A into tensor memory or shared memory
copy_atom_a_transform = self._get_copy_atom_a_transform(
copy_atom_a_transform = mixed_input_utils.get_copy_atom_a_transform(
self.mma_dtype,
self.use_2cta_instrs,
self.transform_a_source,
@@ -928,7 +943,7 @@ class MixedInputGemmKernel:
tCsS = thr_mma.partition_A(sS_input)
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), loopM, loopK, loopL)
tSsS, tSgS = self.scale_tma_partition(
tSsS, tSgS = mixed_input_utils.scale_tma_partition(
tCsS,
tCgS,
tma_atom_s,
@@ -959,10 +974,7 @@ class MixedInputGemmKernel:
)
# Cluster wait before TMEM alloc and ensure pipelines are ready
if cutlass.const_expr(cute.size(self.cluster_shape_mn) > 1):
cute.arch.cluster_wait()
else:
self.cta_sync_barrier.arrive_and_wait()
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
# TMEM allocation
tmem.allocate(self.num_tmem_alloc_cols)
@@ -1145,7 +1157,7 @@ class MixedInputGemmKernel:
dst_copy_a,
tAsA_input,
tAsA_transform,
) = self.transform_partition(
) = mixed_input_utils.transform_partition(
self.transform_a_source,
self.scale_mode,
copy_atom_a_input,
@@ -1173,8 +1185,10 @@ class MixedInputGemmKernel:
tSrS_copy = None
tSrS = None
if cutlass.const_expr(self.scale_mode == TransformMode.ConvertScale):
smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS = self.scale_partition(
src_copy_a, tCsS, transform_local_tidx, self.mma_dtype
smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS = (
mixed_input_utils.scale_partition(
src_copy_a, tCsS, transform_local_tidx, self.mma_dtype
)
)
assert cute.size(tSrS, mode=[0]) == cute.size(tArA, mode=[0]), (
"tSrS and tArA have different leading dimension"
@@ -1582,328 +1596,6 @@ class MixedInputGemmKernel:
if cutlass.const_expr(self.use_tma_store):
c_pipeline.producer_tail()
def scale_tma_partition(
self,
tCsS: cute.Tensor,
tCgS: cute.Tensor,
tma_atom_s: cute.CopyAtom,
block_in_cluster_coord_vmnk: cute.Coord,
scale_cta_layout: cute.Layout,
) -> tuple[cute.Tensor, cute.Tensor]:
"""
Perform TMA partition for scale tensor.
This method partitions the gobal memory and shared memory buffer for scale tensor for TMA load.
:param tCsS: Input scale shared memory tensor
:type tCsS: cute.Tensor
:param tCgS: Input scale global memory tensor
:type tCgS: cute.Tensor
:param tma_atom_s: TMA copy atom for scale tensor
:type tma_atom_s: cute.CopyAtom
:param block_in_cluster_coord_vmnk: CTA coord in the cluster
:type block_in_cluster_coord_vmnk: cute.Coord
:param scale_cta_layout: Layout of CTA from the view of the scale tensor
:type scale_cta_layout: cute.Layout
:return: A tuple containing (tSsS, tSgS) where:
- tSsS: Partitioned scale tensor in shared memory
- tSgS: Partitioned scale tensor in global memory
:rtype: tuple[cute.Tensor, cute.Tensor]
"""
tSsS, tSgS = cpasync.tma_partition(
tma_atom_s,
block_in_cluster_coord_vmnk[2],
scale_cta_layout,
cute.group_modes(tCsS, 0, 3),
cute.group_modes(tCgS, 0, 3),
)
# Add rest_v mode
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), loopM, loopK, loopL)
tSsS = cute.make_tensor(
tSsS.iterator,
cute.make_layout(
((tSsS.layout.shape[0], 1), *tSsS.layout.shape[1:]),
stride=(
(tSsS.layout.stride[0], 0),
*tSsS.layout.stride[1:],
),
),
)
tSgS = cute.make_tensor(
tSgS.iterator,
cute.make_layout(
((tSgS.layout.shape[0], 1), *tSgS.layout.shape[1:]),
stride=(
(tSgS.layout.stride[0], 0),
*tSgS.layout.stride[1:],
),
),
)
return tSsS, tSgS
def transform_partition(
self,
transform_a_source: tcgen05.OperandSource,
scale_mode: TransformMode,
copy_atom_a_input: cute.CopyAtom,
copy_atom_a_transform: cute.CopyAtom,
sA_input: cute.Tensor,
A_transform: cute.Tensor,
transform_local_tidx: cutlass.Int32,
) -> tuple[cute.TiledCopy, cute.TiledCopy, cute.Tensor, cute.Tensor]:
"""
Partition tensors for transform input and output.
This method sets up the copy atoms and partitions the shared/tensor memory
for the transformation of tensor A.
:param transform_a_source: Where the transformed tensor A is stored (TMEM or SMEM)
:type transform_a_source: tcgen05.OperandSource
:param scale_mode: The transform mode (ConvertOnly or ConvertScale)
:type scale_mode: TransformMode
:param copy_atom_a_input: Copy atom for loading A from shared memory
:type copy_atom_a_input: cute.CopyAtom
:param copy_atom_a_transform: Copy atom for storing transformed A
:type copy_atom_a_transform: cute.CopyAtom
:param sA_input: Input tensor A in shared memory
:type sA_input: cute.Tensor
:param A_transform: Transformed tensor A in tensor or shared memory
:type A_transform: cute.Tensor
:param transform_local_tidx: Local thread index for transformation warps
:type transform_local_tidx: cutlass.Int32
:return: A tuple containing (src_copy_a, dst_copy_a, tAsA_input, tA_transform) where:
- src_copy_a: Tiled copy for source tensor
- dst_copy_a: Tiled copy for destination tensor
- tAsA_input: Partitioned input tensor A
- tA_transform: Partitioned transformed tensor A
:rtype: tuple[cute.TiledCopy, cute.TiledCopy, cute.Tensor, cute.Tensor]
"""
if cutlass.const_expr(transform_a_source == tcgen05.OperandSource.TMEM):
if cutlass.const_expr(
cute.size(A_transform, mode=[0, 0]) == 128
and cute.size(sA_input, mode=[0, 0]) == 64
):
tensor_input = cute.make_tensor(
sA_input.iterator,
cute.logical_product(
sA_input.layout,
((cute.make_layout(2, stride=0), None), None, None, None),
),
)
else:
tensor_input = sA_input
reg2tmem_tiled_copy = tcgen05.make_tmem_copy(
copy_atom_a_transform, A_transform[(None, None, None, 0)]
)
thr_reg2tmem_tiled_copy = reg2tmem_tiled_copy.get_slice(
transform_local_tidx
)
partitioned_tensor_input = thr_reg2tmem_tiled_copy.partition_S(tensor_input)
partitioned_tensor_transform = thr_reg2tmem_tiled_copy.partition_D(
A_transform
)
src_copy_a = (
cute.make_tiled_copy_S(copy_atom_a_input, reg2tmem_tiled_copy)
if scale_mode is TransformMode.ConvertScale
else None
)
dst_copy_a = reg2tmem_tiled_copy
tAsA_input = partitioned_tensor_input
tA_transform = partitioned_tensor_transform
elif cutlass.const_expr(transform_a_source == tcgen05.OperandSource.SMEM):
# Construct tiled_copy satisfying 8 contiguous elts per copy atom
reg2smem_tiled_copy = cute.make_cotiled_copy(
copy_atom_a_transform,
cute.make_layout((128, 8), stride=(8, 1)),
A_transform[(None, None, None, 0)].layout,
)
thr_reg2smem_tiled_copy = reg2smem_tiled_copy.get_slice(
transform_local_tidx
)
partitioned_tensor_input = thr_reg2smem_tiled_copy.partition_S(sA_input)
partitioned_tensor_transform = thr_reg2smem_tiled_copy.partition_D(
A_transform
)
src_copy_a = (
cute.make_tiled_copy_S(copy_atom_a_input, reg2smem_tiled_copy)
if scale_mode is TransformMode.ConvertScale
else None
)
# auto-vec copy is enough for copy from register to shared memory here
dst_copy_a = None
tAsA_input = partitioned_tensor_input
tA_transform = partitioned_tensor_transform
return src_copy_a, dst_copy_a, tAsA_input, tA_transform
def scale_partition(
self,
src_copy_a: cute.TiledCopy,
tCsS: cute.Tensor,
transform_local_tidx: cutlass.Int32,
mma_dtype: type[cutlass.Numeric],
) -> tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]:
"""
Partition the scale tensor for transformation.
This method prepares the copy atom and partitions the shared memory for the scale tensor.
:param src_copy_a: Tiled copy for the source tensor
:type src_copy_a: cute.TiledCopy
:param tCsS: Scale tensor in shared memory
:type tCsS: cute.Tensor
:param transform_local_tidx: Local thread index for transformation warps
:type transform_local_tidx: cutlass.Int32
:param mma_dtype: Data type for the MMA operation
:type mma_dtype: type[cutlass.Numeric]
:return: A tuple containing (smem_thr_copy_S, tSsS_trans, tSrS) where:
- smem_thr_copy_S: Tiled copy for the scale tensor
- tSsS_trans: Partitioned scale tensor for transformation
- tSrS_copy: Register fragment for the scale tensor
- tSrS: view of scale tensor used for transformation computation
:rtype: tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]
"""
smem_thr_copy_S = None
tSsS_trans = None
tSrS = None
# Partition scale tensor
smem_thr_copy_S = src_copy_a.get_slice(transform_local_tidx)
tSsS_trans = smem_thr_copy_S.partition_S(tCsS)
# Construct register fragment for scale tensor
tSsS_layout_per_stage = tSsS_trans[(None, None, None, None, 0)].layout
# tSrS for copy
tSrS_copy = cute.make_rmem_tensor(
cute.filter_zeros(tSsS_layout_per_stage).shape, mma_dtype
)
# tSrS view for transformation computation
tSrS = cute.make_tensor(
tSrS_copy.iterator,
cute.make_layout(
tSsS_layout_per_stage.shape, stride=tSrS_copy.layout.stride
),
)
return smem_thr_copy_S, tSsS_trans, tSrS_copy, tSrS
def get_gmem_layout_scale(
self, scale_shape_mkl: tuple[int, int, int]
) -> cute.Layout:
"""
Get the layout of the scale tensor in global memory.
:param scale_shape_mkl: The shape of the scale tensor (M, K, L).
:type scale_shape_mkl: tuple[int, int, int]
:return: The layout of the scale tensor in global memory.
:rtype: cute.Layout
"""
m, k, l = scale_shape_mkl
shape_scale = (
(self.scale_granularity_m, cute.ceil_div(m, self.scale_granularity_m)),
(self.scale_granularity_k, cute.ceil_div(k, self.scale_granularity_k)),
)
if cutlass.const_expr(self.scale_major_mode == tcgen05.OperandMajorMode.MN):
layout_mk = cute.make_layout(
shape_scale,
stride=(
(0, 1),
(0, cute.size(shape_scale[0][1])),
),
)
else:
layout_mk = cute.make_layout(
shape_scale,
stride=(
(0, cute.size(shape_scale[1][1])),
(0, 1),
),
)
return cute.make_layout(
(*layout_mk.shape, l),
stride=(*layout_mk.stride, cute.cosize(layout_mk)),
)
def get_smem_layout_scale(self) -> tuple[cute.ComposedLayout, cute.ComposedLayout]:
"""
Get the layout of the scale tensor in shared memory.
:return: A tuple containing:
- smem_layout_scale_per_stage: Shared memory layout for scale tensor per stage
- smem_layout_scale: Shared memory layout for scale tensor
:rtype: tuple[cute.ComposedLayout, cute.ComposedLayout]
"""
self.scale_tile_shape = (
(
cute.size(self.mma_tiler[0]) // 2
if self.use_2cta_instrs
else cute.size(self.mma_tiler[0])
),
cute.size(self.mma_tiler[2]),
)
size_mn = self.scale_tile_shape[0]
size_k = self.scale_tile_shape[1]
smem_size_mn = (
self.scale_granularity_m if self.scale_granularity_m < size_mn else size_mn
)
smem_size_k = (
self.scale_granularity_k if self.scale_granularity_k < size_k else size_k
)
div_mn = cute.ceil_div(size_mn, smem_size_mn)
div_k = cute.ceil_div(size_k, smem_size_k)
smem_atom_shape = (
(smem_size_mn, div_mn),
(smem_size_k, div_k),
)
if cutlass.const_expr(self.scale_major_mode == tcgen05.OperandMajorMode.MN):
outer_layout = cute.make_layout(
smem_atom_shape,
stride=(
(0, 1),
(0, div_mn),
),
)
else:
outer_layout = cute.make_layout(
smem_atom_shape,
stride=(
(0, div_k),
(0, 1),
),
)
# Apply a trivial swizzle to make it a composed layout, which could be used to construct TMA atom
smem_layout_scale_per_stage = cute.make_composed_layout(
cute.make_swizzle(0, 4, 3), 0, outer_layout
)
assert cute.rank(smem_layout_scale_per_stage) == 2, (
"Scale layout must be rank 2"
)
assert (
cute.size(self.mma_tiler[0])
% cute.size(smem_layout_scale_per_stage.outer[0])
== 0
), "smem_layout_scale_per_stage must equal the tile shape."
assert (
cute.size(self.mma_tiler[2])
% cute.size(smem_layout_scale_per_stage.outer[1])
== 0
), "smem_layout_scale_per_stage must evenly divide tile k shape."
# Shared memory buffer for scale must be at least 128B to satisfy TMA requirement
assert (
cute.size_in_bytes(self.a_scale_dtype, smem_layout_scale_per_stage) >= 128
), "smem size for scale must be at least 128B"
# Scale layout in smem with multiple stages
smem_layout_scale = cute.append(
smem_layout_scale_per_stage,
cute.make_layout(
(self.num_scale_load2trans_stage),
stride=(cute.cosize(smem_layout_scale_per_stage.outer)),
),
)
return smem_layout_scale_per_stage, smem_layout_scale
def epilog_gmem_copy_and_partition(
self,
tidx: cutlass.Int32,
@@ -2285,126 +1977,6 @@ class MixedInputGemmKernel:
num_tmem_a_cols,
)
@staticmethod
def _compute_smem_layout(
tiled_mma: cute.TiledMma,
mma_tiler_mnk: tuple[int, int, int],
a_dtype: type[cutlass.Numeric],
b_dtype: type[cutlass.Numeric],
load2trans_stage_count: int,
trans2mma_stage_count: int,
) -> tuple[
cute.ComposedLayout,
cute.ComposedLayout,
cute.ComposedLayout,
]:
"""
Compute shared memory layouts for tensor A, transformed A and tensor B.
:param tiled_mma: The tiled MMA object defining the core computation.
:type tiled_mma: cute.TiledMma
:param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler.
:type mma_tiler_mnk: tuple[int, int, int]
:param a_dtype: Data type of operand A.
:type a_dtype: type[cutlass.Numeric]
:param b_dtype: Data type of operand B.
:type b_dtype: type[cutlass.Numeric]
:param load2trans_stage_count: Number of stages for load-to-transform pipeline.
:type load2trans_stage_count: int
:param trans2mma_stage_count: Number of stages for transform-to-MMA pipeline.
:type trans2mma_stage_count: int
:return: A tuple containing:
- smem_layout_a: Shared memory layout for tensor A
- smem_layout_a_transform: Shared memory layout for transformed tensor A
- smem_layout_b: Shared memory layout for tensor B
:rtype: tuple[cute.ComposedLayout, cute.ComposedLayout, cute.ComposedLayout]
"""
smem_layout_a = sm100_utils.make_smem_layout_a(
tiled_mma,
mma_tiler_mnk,
a_dtype,
load2trans_stage_count,
)
smem_layout_a_transform = sm100_utils.make_smem_layout_a(
tiled_mma,
mma_tiler_mnk,
tiled_mma.op.a_dtype,
trans2mma_stage_count,
)
smem_layout_b = sm100_utils.make_smem_layout_b(
tiled_mma,
mma_tiler_mnk,
b_dtype,
load2trans_stage_count,
)
return (
smem_layout_a,
smem_layout_a_transform,
smem_layout_b,
)
@staticmethod
def _get_transform_a_source(
a_major_mode: tcgen05.OperandMajorMode,
) -> tcgen05.OperandSource:
"""
Determine the operand source for transformed A tensor based on the operand major mode.
"""
if cutlass.const_expr(a_major_mode == tcgen05.OperandMajorMode.K):
return tcgen05.OperandSource.TMEM
else:
return tcgen05.OperandSource.SMEM
@staticmethod
def _get_tma_atom_kind(
mcast: cutlass.Boolean,
use_2cta_instrs: bool,
is_b: bool,
) -> Union[
cpasync.CopyBulkTensorTileG2SMulticastOp, cpasync.CopyBulkTensorTileG2SOp
]:
"""
Get the TMA atom kind based on 1) whether it's a multicast operation,
2) whether 2CTA tcgen05.mma instruction is enabled, and
3) whether it's a B tensor
"""
# Not using .2CTA instructions for tensor A as the consumer is threads on different CTAs
cta_group = (
tcgen05.CtaGroup.TWO if (use_2cta_instrs and is_b) else tcgen05.CtaGroup.ONE
)
if cutlass.const_expr(mcast):
return cpasync.CopyBulkTensorTileG2SMulticastOp(cta_group)
return cpasync.CopyBulkTensorTileG2SOp(cta_group)
@staticmethod
def _get_copy_atom_a_transform(
mma_dtype: type[cutlass.Numeric],
use_2cta_instrs: bool,
transform_a_source: tcgen05.OperandSource,
a_smem_shape: cute.Shape,
a_dtype: type[cutlass.Numeric],
) -> cute.CopyAtom:
"""
Determine the copy atom for transformed A tensor based on the operand source and tile size.
"""
if cutlass.const_expr(transform_a_source == tcgen05.OperandSource.TMEM):
if cutlass.const_expr(
cute.size(a_smem_shape[0][0]) == 64 and (not use_2cta_instrs)
):
copy_op_r2t = tcgen05.St16x256bOp(
tcgen05.Repetition(1), tcgen05.Unpack.NONE
)
else:
copy_op_r2t = tcgen05.St32x32bOp(
tcgen05.Repetition(8), tcgen05.Unpack.NONE
)
return cute.make_copy_atom(copy_op_r2t, mma_dtype)
else:
return cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(), a_dtype, num_bits_per_copy=32
)
@staticmethod
def _compute_grid(
c: cute.Tensor,
@@ -2429,29 +2001,6 @@ class MixedInputGemmKernel:
return tile_sched_params, grid
def is_valid_scale_granularity(
scale_granularity_m: int,
scale_granularity_k: int,
a_dtype: type[cutlass.Numeric],
k: int,
mma_tiler_k: int,
) -> bool:
"""
Check if the scale granularity settings are valid for the given data type and problem size.
"""
if a_dtype.width == 8:
# No scale tensor for 8bit data type A
if not (scale_granularity_m == 0 and scale_granularity_k == 0):
return False
elif a_dtype.width == 4:
if scale_granularity_m != 1 or (
scale_granularity_k == 0
or k % scale_granularity_k != 0
or scale_granularity_k % mma_tiler_k != 0
):
return False
return True
def is_valid_tensor_alignment(
m: int,
n: int,
@@ -2566,7 +2115,7 @@ class MixedInputGemmKernel:
mma_tiler, cluster_shape_mn, use_2cta_instrs
):
return False
if not MixedInputGemmKernel.is_valid_scale_granularity(
if not mixed_input_utils.is_valid_scale_granularity(
scale_granularity_m, scale_granularity_k, a_dtype, k, mma_tiler[2]
):
return False
@@ -2634,7 +2183,11 @@ def create_i4_tensor_and_scale(
m, num_scales, scale_granularity_k, l
)
# Get elements with maximum absolute value to compute scaling factors
a_max = torch.maximum(ref / up_4b, ref / lb_4b)
a_max = (
torch.maximum(ref / up_4b, ref / lb_4b)
if dtype == cutlass.Int4
else torch.maximum(ref / up_4b)
)
a_scales, _ = torch.max(a_max, dim=2, keepdim=True)
a_scale_inv = torch.where(a_scales == 0, 0, 1 / a_scales)
a_quant = ref * a_scale_inv
@@ -2668,17 +2221,6 @@ def create_i4_tensor_and_scale(
)
def get_divisibility(contiguous_dim_size: int, upper_bound: int = 128) -> int:
"""
Calculate the largest power of 2 divisibility factor for memory alignment.
"""
# Check the largest power of 2 factor of contiguous_dim_size
for i in range(int(log2(contiguous_dim_size)), 0, -1):
if contiguous_dim_size % (2**i) == 0:
return min(2**i, upper_bound)
return 1
def create_tensor_a(
l: int,
m: int,
@@ -2688,7 +2230,7 @@ def create_tensor_a(
scale_granularity_m: int = 0,
scale_granularity_k: int = 0,
transformed_dtype: Optional[type[cutlass.Numeric]] = None,
) -> tuple[cute.Tensor, cute.Tensor, torch.Tensor, torch.Tensor]:
) -> tuple[cute.Tensor, Optional[cute.Tensor], torch.Tensor, Optional[torch.Tensor]]:
"""
Create tensor A and scale tensor.
"""
@@ -2710,7 +2252,7 @@ def create_tensor_a(
a_dtype,
scale_granularity_m,
scale_granularity_k,
divisibility=get_divisibility(m if a_major == "m" else k),
divisibility=mixed_input_utils.get_divisibility(m if a_major == "m" else k),
transformed_dtype=transformed_dtype,
)
else:
@@ -2725,7 +2267,9 @@ def create_tensor_a(
a_torch_cpu,
a_dtype,
is_dynamic_layout=True,
assumed_align=get_divisibility(m if a_major == "m" else k),
assumed_align=mixed_input_utils.get_divisibility(
m if a_major == "m" else k
),
)
return a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu
@@ -2774,18 +2318,18 @@ def create_tensors(
b_torch_cpu,
b_dtype,
is_dynamic_layout=True,
assumed_align=get_divisibility(n if b_major == "n" else k),
assumed_align=mixed_input_utils.get_divisibility(n if b_major == "n" else k),
)
c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like(
c_torch_cpu,
c_dtype,
is_dynamic_layout=True,
assumed_align=get_divisibility(m if c_major == "m" else n),
assumed_align=mixed_input_utils.get_divisibility(m if c_major == "m" else n),
)
c_tensor = c_tensor.mark_compact_shape_dynamic(
mode=(0 if c_major == "m" else 1),
stride_order=(2, 1, 0) if c_major == "m" else (2, 0, 1),
divisibility=get_divisibility(m if c_major == "m" else n),
divisibility=mixed_input_utils.get_divisibility(m if c_major == "m" else n),
)
return (
@@ -2966,25 +2510,36 @@ def run(
def generate_tensors():
a_tensor, a_scale_tensor, a_torch_cpu, a_scale_torch_cpu = create_tensor_a(
l, m, k, a_major, a_dtype, scale_granularity_m, scale_granularity_k, b_dtype
l,
m,
k,
a_major,
a_dtype,
scale_granularity_m,
scale_granularity_k,
b_dtype,
)
b_tensor, _ = cutlass_torch.cute_tensor_like(
b_torch_cpu,
b_dtype,
is_dynamic_layout=True,
assumed_align=get_divisibility(n if b_major == "n" else k),
assumed_align=mixed_input_utils.get_divisibility(
n if b_major == "n" else k
),
)
c_torch_cpu = cutlass_torch.matrix(l, m, n, c_major == "m", c_dtype)
c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like(
c_torch_cpu,
c_dtype,
is_dynamic_layout=True,
assumed_align=get_divisibility(m if c_major == "m" else n),
assumed_align=mixed_input_utils.get_divisibility(
m if c_major == "m" else n
),
)
c_tensor = c_tensor.mark_compact_shape_dynamic(
mode=(0 if c_major == "m" else 1),
stride_order=(2, 1, 0) if c_major == "m" else (2, 0, 1),
divisibility=get_divisibility(m if c_major == "m" else n),
divisibility=mixed_input_utils.get_divisibility(m if c_major == "m" else n),
)
return testing.JitArguments(
a_tensor, a_scale_tensor, b_tensor, c_tensor, current_stream
examples/python/CuTeDSL/blackwell/mla.py
+13 -9
View File
@@ -43,6 +43,7 @@ import cutlass.cute.nvgpu.tcgen05 as tcgen05
import cutlass.cute.nvgpu.cpasync as cpasync
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
import cutlass.torch as cutlass_torch
import cutlass.utils.blackwell_helpers as sm100_utils
from cutlass.cute.runtime import from_dlpack
@@ -970,7 +971,6 @@ class BlackwellMultiHeadLatentAttentionForward:
num_tmem_dealloc_threads = self.threads_per_warp * self.num_compute_warps
with cute.arch.elect_one():
cute.arch.mbarrier_init(tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads)
cute.arch.mbarrier_init_fence()
load_q_pipeline = self.make_and_init_load_qkv_pipeline(
storage.load_q_mbar_ptr.data_ptr(),
@@ -1004,8 +1004,7 @@ class BlackwellMultiHeadLatentAttentionForward:
)
# Cluster arrive after barrier init
if cutlass.const_expr(cute.size(self.cluster_shape_mnk) > 1):
cute.arch.cluster_arrive_relaxed()
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True)
# Generate smem tensor Q/KC/VC/exchange
# (MMA, MMA_H, MMA_R, PIPE)
@@ -1035,10 +1034,7 @@ class BlackwellMultiHeadLatentAttentionForward:
#
# Cluster wait before tensor memory alloc
#
if cutlass.const_expr(cute.size(self.cluster_shape_mnk) > 1):
cute.arch.cluster_wait()
else:
pipeline.sync(barrier_id=4)
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk)
# ///////////////////////////////////////////////////////////////////////////////
# Load warps, including page table and data tensors
@@ -2046,8 +2042,9 @@ class BlackwellMultiHeadLatentAttentionForward:
# wait cpasync arrive until the last stage
load_kv_pipeline = common_params.load_kv_pipeline
if copy_in_flight_count == self.load_kv_stage:
cute.arch.cp_async_wait_group(self.load_kv_stage - 1)
release_distance = 2
if copy_in_flight_count == self.load_kv_stage - release_distance:
cute.arch.cp_async_wait_group(self.load_kv_stage - release_distance - 1)
load_kv_pipeline.producer_commit(load_kv_commit_state)
load_kv_commit_state.advance()
copy_in_flight_count -= 1
@@ -3922,6 +3919,7 @@ class BlackwellMultiHeadLatentAttentionForward:
producer_group=load_qkv_producer_group,
consumer_group=load_qkv_consumer_group,
cta_layout_vmnk=cta_layout_vmnk,
defer_sync=True,
)
else:
load_qkv_producer_group = pipeline.CooperativeGroup(
@@ -3937,6 +3935,7 @@ class BlackwellMultiHeadLatentAttentionForward:
consumer_group=load_qkv_consumer_group,
tx_count=tx_count,
cta_layout_vmnk=cta_layout_vmnk,
defer_sync=True,
)
def make_and_init_mma_s_pipeline(
@@ -3971,6 +3970,7 @@ class BlackwellMultiHeadLatentAttentionForward:
producer_group=mma_s_producer_group,
consumer_group=mma_s_consumer_group,
cta_layout_vmnk=cta_layout_vmnk,
defer_sync=True,
)
def make_and_init_p_mma_pipeline(
@@ -4005,6 +4005,7 @@ class BlackwellMultiHeadLatentAttentionForward:
producer_group=p_mma_producer_group,
consumer_group=p_mma_consumer_group,
cta_layout_vmnk=cta_layout_vmnk,
defer_sync=True,
)
def make_and_init_p_cor_pipeline(
@@ -4033,6 +4034,7 @@ class BlackwellMultiHeadLatentAttentionForward:
num_stages=self.p_cor_stage,
producer_group=p_cor_producer_group,
consumer_group=p_cor_consumer_group,
defer_sync=True,
)
def make_and_init_mma_o_pipeline(
@@ -4067,6 +4069,7 @@ class BlackwellMultiHeadLatentAttentionForward:
producer_group=mma_o_producer_group,
consumer_group=mma_o_consumer_group,
cta_layout_vmnk=cta_layout_vmnk,
defer_sync=True,
)
def make_and_init_load_pt_pipeline(self, load_pt_mbar_ptr):
@@ -4091,6 +4094,7 @@ class BlackwellMultiHeadLatentAttentionForward:
num_stages=self.load_pt_stage,
producer_group=load_pt_producer_group,
consumer_group=load_pt_consumer_group,
defer_sync=True,
)
@staticmethod