submission 480794
John · python · License unknown
Use it
Vendorable · source mirrored · license unknownView source →
No package. Vendor the mirrored source: 2475 lines, June 9 Researcher Reciprocity License v1.0.
submission_manual.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-nvfp4-group-gemm-480794?include=source"interfacepython
Compatibility
measured onNVIDIA B200
declared hardwareNVIDIA B200
architecturessm_100
dtypesfp8_e4m3, nvfp4
Benchmark evidence
1 measurement across 1 GPU, fastest first.
Operation / workload
Hardware
Latency
Rank
Observed
Reported · How evidence levels are derived →
Source and license
sourceavailable
revision digestsha256:b5147a420dedb9786101bff9983d785a7e523a4ce0f43000a140f5a3068f873d
license declaredunknown
license concludedunknown
authorsJohn
imported2026-08-15
Techniques
Extracted from the mirrored source by pattern, never inferred. Each row cites its line.
fused-epilogue
epi_tile = sm100_utils.compute_epilogue_tile_shape(mbarrier
self.epilog_sync_barrier = pipeline.NamedBarrier(persistent-kernel
tile_sched_params: utils.PersistentTileSchedulerParams,shared-memory
smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")tcgen05
tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONEwarp-specialization
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)Kernel source
submission_manual.py2475 lines
#!POPCORN leaderboard nvfp4_group_gemm
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.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
from cutlass.cute.runtime import make_ptr
from inspect import isclass
import functools
from typing import Tuple, List, Type, Union
import torch
from task import input_t, output_t
# Kernel configuration parameters
# Size of tma descriptor in bytes
bytes_per_tensormap = 128
# Number of tensormaps: a, b, sfa, sfb
num_tensormaps = 4
# Tile sizes for M, N, K dimensions
mma_tiler_mnk = (128, 128, 256)
# Shape of the K dimension for the MMA instruction
mma_inst_shape_k = 64
mma_inst_tile_k = 4
# FP4 data type for A and B
ab_dtype = cutlass.Float4E2M1FN
# FP8 data type for scale factors
sf_dtype = cutlass.Float8E4M3FN
# FP16 output type
c_dtype = cutlass.Float16
# Scale factor block size (16 elements share one scale)
sf_vec_size = 16
# Total number of columns in tmem
num_tmem_alloc_cols = 512
smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
buffer_align_bytes = 1024
"""
A sensible default is (1, 1)—that means no CTA clustering. Each cluster is just one CTA, so:
- simplest scheduling,
- no multicast,
- lowest coordination overhead,
- good baseline for correctness and first‑time tuning.
You typically consider larger cluster shapes when:
- M and N are both large (more parallelism),
- you want multicast reuse of A/B across CTAs,
- and you’re targeting performance on SM100.
"""
cluster_shape_mn = (1, 2)
class Sm100GroupedBlockScaledGemmKernel:
# Size of smem we reserved for mbarrier, tensor memory management and tensormap update
reserved_smem_bytes = 1024
# bytes_per_tensormap is the size of one TMA descriptor (tensormap).
# In this kernel it’s fixed at 128 bytes, which is what the hardware expects for a TMA descriptor on SM100.
bytes_per_tensormap = 128
# 5 TMA descriptor
# A, B, sfA, sfB, C
num_tensormaps = 5
tensor_memory_management_bytes = 12
def __init__(
self,
mma_tiler_mn: Tuple[int, int],
cluster_shape_mn: Tuple[int, int],
):
self.acc_dtype = cutlass.Float32
self.use_2cta_instrs = mma_tiler_mn[0] == 256
self.cluster_shape_mn = cluster_shape_mn
# K dimension is deferred in _setup_attributes
self.mma_tiler = (*mma_tiler_mn, 1)
self.cta_group = (
tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
)
self.tensormap_update_mode = utils.TensorMapUpdateMode.SMEM
self.epilog_warp_id = (
0,
1,
2,
3,
)
self.mma_warp_id = 4
self.tma_warp_id = 5
self.occupancy = 1
self.threads_per_cta = 32 * len(
(self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
)
# bytes_per_tensormap is in bytes, so dividing by 8 converts bytes → number of int64 elements
self.size_tensormap_in_i64 = (
Sm100GroupedBlockScaledGemmKernel.num_tensormaps
* Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap
// 8
)
# Set barrier for epilogue sync and tmem ptr sync
self.epilog_sync_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=32 * len(self.epilog_warp_id),
)
self.tmem_alloc_barrier = pipeline.NamedBarrier(
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=3,
num_threads=64,
)
@cute.kernel
def ported_kernel(
self,
tiled_mma: cute.TiledMma,
tiled_mma_sfb: cute.TiledMma,
tma_atom_a: cute.CopyAtom,
mA_mkl: cute.Tensor,
tma_atom_b: cute.CopyAtom,
mB_nkl: cute.Tensor,
tma_atom_sfa: cute.CopyAtom,
mSFA_mkl: cute.Tensor,
tma_atom_sfb: cute.CopyAtom,
mSFB_nkl: cute.Tensor,
tma_atom_c: cute.CopyAtom,
mC_mnl: cute.Tensor,
cluster_layout_vmnk: cute.Layout,
cluster_layout_sfb_vmnk: cute.Layout,
a_smem_layout_staged: cute.ComposedLayout,
b_smem_layout_staged: cute.ComposedLayout,
sfa_smem_layout_staged: cute.Layout,
sfb_smem_layout_staged: cute.Layout,
c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout],
epi_tile: cute.Tile,
tile_sched_params: utils.PersistentTileSchedulerParams,
group_count: cutlass.Constexpr,
problem_sizes_mnkl: cute.Tensor,
ptrs_abc: cute.Tensor,
ptrs_sfasfb: cute.Tensor,
tensormaps: cute.Tensor,
strides_abc: cute.Tensor
):
warp_idx = cute.arch.warp_idx()
# Hinting to the compiler this is identical across warp, enables warp level control flow, no divergent branches, better codegen.
warp_idx = cute.arch.make_warp_uniform(warp_idx)
# It runs only in the dedicated TMA warp and prefetches the TMA descriptors for A/B/SFA/SFB/C into the hardware’s descriptor cache.
# That reduces latency when the warp later issues TMA loads/stores using those descriptors.
# This is just warming the cache; it doesn’t move any actual tensor data yet.
if warp_idx == self.tma_warp_id:
cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_a)
cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_b)
cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_sfa)
cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_sfb)
cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_c)
use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2
# Figuring out where the current CTA and thread sit inside the cluster so it can pick the right slice of work.
bidx, bidy, bidz = cute.arch.block_idx()
mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
is_leader_cta = mma_tile_coord_v == 0
# Noting, best practice to hint warp level variables.
cta_rank_in_cluster = cute.arch.make_warp_uniform(
cute.arch.block_idx_in_cluster()
)
block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(
cta_rank_in_cluster
)
block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(
cta_rank_in_cluster
)
# coord inside cta
tidx, _, _ = cute.arch.thread_idx()
# Shared memory is just a raw byte buffer per CTA.
# You need a structured layout to carve it into named regions (tensormap buffer, mbarriers, sA/sB/sC tiles, etc.).
# @cute.struct lets you describe that layout in one place with sizes and alignment
# so later code can refer to storage.sA, storage.sB, etc. without manual pointer math.
smem = utils.SmemAllocator()
storage = smem.allocate(self.shared_storage)
tensormap_smem_ptr = storage.tensormap_buffer.data_ptr()
tensormap_a_smem_ptr = tensormap_smem_ptr
tensormap_b_smem_ptr = (
tensormap_a_smem_ptr
+ Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8
)
tensormap_sfa_smem_ptr = (
tensormap_b_smem_ptr
+ Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8
)
tensormap_sfb_smem_ptr = (
tensormap_sfa_smem_ptr
+ Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8
)
tensormap_c_smem_ptr = (
tensormap_sfb_smem_ptr
+ Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8
)
tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar_ptr
tmem_holding_buf = storage.tmem_holding_buf
# Initialize mainloop ab_pipeline (barrier) and states
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
ab_pipeline_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, num_tma_producer
)
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,
consumer_group=ab_pipeline_consumer_group,
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cluster_layout_vmnk,
)
# Initialize acc_pipeline (barrier) and states
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 = pipeline.CooperativeGroup(
pipeline.Agent.Thread, num_acc_consumer_threads
)
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,
)
# Tensor memory dealloc barrier init
if use_2cta_instrs:
if warp_idx == self.tma_warp_id:
num_tmem_dealloc_threads = 32
with cute.arch.elect_one():
cute.arch.mbarrier_init(
tmem_dealloc_mbar_ptr, num_tmem_dealloc_threads
)
# Cluster arrive after barrier init
# The pipeline uses mbarriers to coordinate between CTAs in a cluster.
# At startup, those barriers are in an “empty/not arrived” state. pipeline_init_arrive(...) marks the initial arrival
# so the other CTAs (or later code) that do a pipeline_init_wait(...) don’t deadlock on a barrier
# that was never initialized/signaled. It’s basically “prime the barrier state” for cluster‑wide use.
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
#
# Setup smem tensor A/B/SFA/SFB/C
# The swizzle are detemrined by a blackwell helper
# it’s chosen to make the major dimension align to 128/64/32 bytes (1024/512/256 bits), which is optimal for TMA and SMEM access
#
sC = storage.sC.get_tensor(
c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner
)
# (MMA, MMA_M, MMA_K, STAGE)
sA = storage.sA.get_tensor(
a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner
)
# (MMA, MMA_N, MMA_K, STAGE)
sB = storage.sB.get_tensor(
b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner
)
# (MMA, MMA_M, MMA_K, STAGE)
sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
# (MMA, MMA_N, MMA_K, STAGE)
sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged)
#
# Compute multicast mask for A/B/SFA/SFB buffer full
#
a_full_mcast_mask = None
b_full_mcast_mask = None
sfa_full_mcast_mask = None
sfb_full_mcast_mask = None
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
)
b_full_mcast_mask = cpasync.create_tma_multicast_mask(
cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
)
sfa_full_mcast_mask = cpasync.create_tma_multicast_mask(
cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
)
sfb_full_mcast_mask = cpasync.create_tma_multicast_mask(
cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1
)
#
# Local_tile partition global tensors
# Level = (which tile of A/B/C this CTA works on
#
# (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, RestN, RestK, RestL)
gB_nkl = cute.local_tile(
mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
# (bM, bK, RestM, RestK, RestL)
gSFA_mkl = cute.local_tile(
mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
# (bN, bK, RestN, RestK, RestL)
gSFB_nkl = cute.local_tile(
mSFB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
# (bM, bN, RestM, RestN, RestL)
gC_mnl = cute.local_tile(
mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
)
#
# Partition global tensor for TiledMMA_A/B/C
# Level = what each warp/lane loads/computes
#
thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v)
# (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
tCgA = thr_mma.partition_A(gA_mkl)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgB = thr_mma.partition_B(gB_nkl)
# (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
tCgSFA = thr_mma.partition_A(gSFA_mkl)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl)
# (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
tCgC = thr_mma.partition_C(gC_mnl)
"""
Yes—TMA also needs partitioning, but for a different reason than MMA.
Why partition for TMA:
- TMA moves tiles from GMEM → SMEM.
- In clustered CTAs, different CTAs are responsible for different slices of those tiles (and some are multicast).
- cpasync.tma_partition(...) computes the per‑CTA source (tAgA/tBgB) and destination (tAsA/tBsB) slices for the TMA operation.
So:
- MMA partitioning is within a CTA (per‑warp/per‑thread).
- TMA partitioning is across CTAs in a cluster, deciding which CTA loads which part and where it lands in SMEM.
"""
#
# Partition global/shared tensor for TMA load A/B
#
# TMA load A partition_S/D
a_cta_layout = cute.make_layout(
cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
)
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestM, RestK, RestL)
tAsA, tAgA = cpasync.tma_partition(
tma_atom_a,
block_in_cluster_coord_vmnk[2],
a_cta_layout,
cute.group_modes(sA, 0, 3),
cute.group_modes(tCgA, 0, 3),
)
# TMA load B partition_S/D
b_cta_layout = cute.make_layout(
cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
)
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestN, RestK, RestL)
tBsB, tBgB = cpasync.tma_partition(
tma_atom_b,
block_in_cluster_coord_vmnk[1],
b_cta_layout,
cute.group_modes(sB, 0, 3),
cute.group_modes(tCgB, 0, 3),
)
# TMA Load SFA partition_S/D
sfa_cta_layout = a_cta_layout
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestM, RestK, RestL)
tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition(
tma_atom_sfa,
block_in_cluster_coord_vmnk[2],
sfa_cta_layout,
cute.group_modes(sSFA, 0, 3),
cute.group_modes(tCgSFA, 0, 3),
)
tAsSFA = cute.filter_zeros(tAsSFA)
tAgSFA = cute.filter_zeros(tAgSFA)
# TMA Load SFB partition_S/D
sfb_cta_layout = cute.make_layout(
cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape
)
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestN, RestK, RestL)
tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition(
tma_atom_sfb,
block_in_cluster_coord_sfb_vmnk[1],
sfb_cta_layout,
cute.group_modes(sSFB, 0, 3),
cute.group_modes(tCgSFB, 0, 3),
)
tBsSFB = cute.filter_zeros(tBsSFB)
tBgSFB = cute.filter_zeros(tBgSFB)
#
# Partition shared/tensor memory tensor for TiledMMA_A/B/C
#
# (MMA, MMA_M, MMA_K, STAGE)
tCrA = tiled_mma.make_fragment_A(sA)
# (MMA, MMA_N, MMA_K, STAGE)
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)
)
#
# Cluster wait before tensor memory alloc
# pipeline_init_wait(...) is the cluster‑wide “wait for initialization” barrier.
# It makes each CTA in the cluster wait until all CTAs have finished pipeline_init_arrive(...) (i.e., the pipeline barriers are initialized and in a consistent state).
# This prevents any CTA from using the pipeline/barriers before others have finished setting them up, which avoids deadlocks or incorrect synchronization.
#
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
#
# Get tensormap buffer address
# compute the pointer(s) to the TMA descriptor storage for this CTA’s workspace, so the kernel can read/update those descriptors.
#
grid_dim = cute.arch.grid_dim()
tensormap_workspace_idx = (
bidz * grid_dim[1] * grid_dim[0] + bidy * grid_dim[0] + bidx
)
# The host allocates tensormap_cute_tensor, but the kernel creates/updates the actual descriptor contents at runtime.
# tensormap_cute_tensor is the global-memory backing store for all TMA descriptors.
# It’s passed into the kernel so each CTA can compute its slice of that tensor and then read/update the descriptors for A/B/SF/C at runtime.
tensormap_manager = utils.TensorMapManager(
utils.TensorMapUpdateMode.SMEM,
Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap,
)
tensormap_a_gmem_ptr = tensormap_manager.get_tensormap_ptr(
tensormaps[(tensormap_workspace_idx, 0, None)].iterator
)
tensormap_b_gmem_ptr = tensormap_manager.get_tensormap_ptr(
tensormaps[(tensormap_workspace_idx, 1, None)].iterator
)
tensormap_sfa_gmem_ptr = tensormap_manager.get_tensormap_ptr(
tensormaps[(tensormap_workspace_idx, 2, None)].iterator
)
tensormap_sfb_gmem_ptr = tensormap_manager.get_tensormap_ptr(
tensormaps[(tensormap_workspace_idx, 3, None)].iterator
)
tensormap_c_gmem_ptr = tensormap_manager.get_tensormap_ptr(
tensormaps[(tensormap_workspace_idx, 4, None)].iterator
)
#
# Specialized TMA load warp
#
if warp_idx == self.tma_warp_id:
#
# Persistent tile scheduling loop
#
tile_sched = utils.StaticPersistentTileScheduler.create(
tile_sched_params, cute.arch.block_idx(), grid_dim
)
# grouped gemm tile scheduler helper will compute the group index for the tile we're working on
group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper(
group_count,
tile_sched_params,
self.cluster_tile_shape_mnk,
utils.create_initial_search_state(),
)
tensormap_init_done = cutlass.Boolean(False)
# group index of last tile
last_group_idx = cutlass.Int32(-1)
work_tile = tile_sched.initial_work_tile_info()
# If producer and consumer use different num_stages, they’ll index different ring sizes and the barrier protocol breaks (stages will be marked full/empty out of sync).
# In practice, that would deadlock or corrupt data. The pipeline API assumes both sides are created with the same stage count.
ab_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_ab_stage
)
while work_tile.is_valid_tile:
cur_tile_coord = work_tile.tile_idx
grouped_gemm_cta_tile_info = group_gemm_ts_helper.delinearize_z(
cur_tile_coord,
problem_sizes_mnkl,
)
cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k
cur_group_idx = grouped_gemm_cta_tile_info.group_idx
is_group_changed = cur_group_idx != last_group_idx
# skip tensormap update if we're working on the same group
if is_group_changed:
real_tensor_a = self.make_tensor_abc_for_tensormap_update(
cur_group_idx,
ab_dtype,
(
grouped_gemm_cta_tile_info.problem_shape_m,
grouped_gemm_cta_tile_info.problem_shape_n,
grouped_gemm_cta_tile_info.problem_shape_k,
),
strides_abc,
ptrs_abc,
0, # 0 for tensor A
)
real_tensor_b = self.make_tensor_abc_for_tensormap_update(
cur_group_idx,
ab_dtype,
(
grouped_gemm_cta_tile_info.problem_shape_m,
grouped_gemm_cta_tile_info.problem_shape_n,
grouped_gemm_cta_tile_info.problem_shape_k,
),
strides_abc,
ptrs_abc,
1, # 1 for tensor B
)
real_tensor_sfa = self.make_tensor_sfasfb_for_tensormap_update(
cur_group_idx,
sf_dtype,
(
grouped_gemm_cta_tile_info.problem_shape_m,
grouped_gemm_cta_tile_info.problem_shape_n,
grouped_gemm_cta_tile_info.problem_shape_k,
),
ptrs_sfasfb,
0, # 0 for tensor SFA
)
real_tensor_sfb = self.make_tensor_sfasfb_for_tensormap_update(
cur_group_idx,
sf_dtype,
(
grouped_gemm_cta_tile_info.problem_shape_m,
grouped_gemm_cta_tile_info.problem_shape_n,
grouped_gemm_cta_tile_info.problem_shape_k,
),
ptrs_sfasfb,
1, # 1 for tensor SFB
)
if tensormap_init_done == False:
# wait tensormap initialization complete
self.tensormap_ab_init_barrier.arrive_and_wait()
tensormap_init_done = True
tensormap_manager.update_tensormap(
(
real_tensor_a,
real_tensor_b,
real_tensor_sfa,
real_tensor_sfb,
),
(tma_atom_a, tma_atom_b, tma_atom_sfa, tma_atom_sfb),
(
tensormap_a_gmem_ptr,
tensormap_b_gmem_ptr,
tensormap_sfa_gmem_ptr,
tensormap_sfb_gmem_ptr,
),
self.tma_warp_id,
(
tensormap_a_smem_ptr,
tensormap_b_smem_ptr,
tensormap_sfa_smem_ptr,
tensormap_sfb_smem_ptr,
),
)
# M, N, L coordinates of the MMA tile within the current CTA tile
mma_tile_coord_mnl = (
grouped_gemm_cta_tile_info.cta_tile_idx_m
// cute.size(tiled_mma.thr_id.shape),
grouped_gemm_cta_tile_info.cta_tile_idx_n,
0,
)
#
# Slice to per mma tile index
#
# ((atom_v, rest_v), RestK)
tAgA_slice = tAgA[
(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])
]
# ((atom_v, rest_v), RestK)
tBgB_slice = tBgB[
(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])
]
# ((atom_v, rest_v), RestK)
tAgSFA_slice = tAgSFA[
(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])
]
# ((atom_v, rest_v), RestK)
tBgSFB_slice = tBgSFB[
(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])
]
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt
#
# Why: If you skip the peek and just call producer_acquire, you may stall immediately even when there’s no work ready to consume yet.
# The peek lets the warp prefetch optimistically and only block if the stage is actually full. Without it, you can add unnecessary bubbles in the TMA pipeline.
ab_producer_state.reset_count()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < cur_k_tile_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(
ab_producer_state
)
if is_group_changed:
tensormap_manager.fence_tensormap_update(tensormap_a_gmem_ptr)
tensormap_manager.fence_tensormap_update(tensormap_b_gmem_ptr)
tensormap_manager.fence_tensormap_update(tensormap_sfa_gmem_ptr)
tensormap_manager.fence_tensormap_update(tensormap_sfb_gmem_ptr)
tma_desc_ptr_a = tensormap_manager.get_tensormap_ptr(
tensormap_a_gmem_ptr,
cute.AddressSpace.generic,
)
tma_desc_ptr_b = tensormap_manager.get_tensormap_ptr(
tensormap_b_gmem_ptr,
cute.AddressSpace.generic,
)
tma_desc_ptr_sfa = tensormap_manager.get_tensormap_ptr(
tensormap_sfa_gmem_ptr,
cute.AddressSpace.generic,
)
tma_desc_ptr_sfb = tensormap_manager.get_tensormap_ptr(
tensormap_sfb_gmem_ptr,
cute.AddressSpace.generic,
)
#
# Tma load loop
#
for k_tile in cutlass.range(0, cur_k_tile_cnt, 1, unroll=1):
# Conditionally wait for AB buffer empty
ab_pipeline.producer_acquire(
ab_producer_state, peek_ab_empty_status
)
# Issue TMA load A/B/SFA/SFB
cute.copy(
tma_atom_a,
tAgA_slice[(None, ab_producer_state.count)],
tAsA[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=a_full_mcast_mask,
tma_desc_ptr=tma_desc_ptr_a,
)
cute.copy(
tma_atom_b,
tBgB_slice[(None, ab_producer_state.count)],
tBsB[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=b_full_mcast_mask,
tma_desc_ptr=tma_desc_ptr_b,
)
cute.copy(
tma_atom_sfa,
tAgSFA_slice[(None, ab_producer_state.count)],
tAsSFA[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=sfa_full_mcast_mask,
tma_desc_ptr=tma_desc_ptr_sfa,
)
cute.copy(
tma_atom_sfb,
tBgSFB_slice[(None, ab_producer_state.count)],
tBsSFB[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=sfb_full_mcast_mask,
tma_desc_ptr=tma_desc_ptr_sfb,
)
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1
ab_producer_state.advance()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < cur_k_tile_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(
ab_producer_state
)
#
# Advance to next tile
#
tile_sched.advance_to_next_work()
work_tile = tile_sched.get_current_work()
last_group_idx = cur_group_idx
#
# Wait A/B buffer empty
#
ab_pipeline.producer_tail(ab_producer_state)
#
# Specialized MMA warp
#
if warp_idx == self.mma_warp_id:
#
# Initialize tensormaps for A, B, SFA and SFB
#
# Not specifically used in the MMA warp, one of the warp just had to initialise it and create
# the base descriptior in the shared memory.
#
tensormap_manager.init_tensormap_from_atom(
tma_atom_a, tensormap_a_smem_ptr, self.mma_warp_id
)
tensormap_manager.init_tensormap_from_atom(
tma_atom_b, tensormap_b_smem_ptr, self.mma_warp_id
)
tensormap_manager.init_tensormap_from_atom(
tma_atom_sfa, tensormap_sfa_smem_ptr, self.mma_warp_id
)
tensormap_manager.init_tensormap_from_atom(
tma_atom_sfb, tensormap_sfb_smem_ptr, self.mma_warp_id
)
# indicate tensormap initialization has finished
self.tensormap_ab_init_barrier.arrive_and_wait()
#
# Bar sync for retrieve tensor memory ptr from shared mem
#
self.tmem_alloc_barrier.arrive_and_wait()
#
# Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor
#
# Make accumulator tmem tensor
acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
self.acc_dtype,
alignment=16,
ptr_to_buffer_holding_addr=tmem_holding_buf,
)
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
# Make SFA tmem tensor
sfa_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc_base),
dtype=sf_dtype,
)
# (MMA, MMA_M, MMA_K)
tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
tiled_mma,
self.mma_tiler,
sf_vec_size,
cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
)
tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
# 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),
dtype=sf_dtype,
)
# (MMA, MMA_N, MMA_K)
tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
tiled_mma,
self.mma_tiler,
sf_vec_size,
cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
)
tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout)
#
# Partition for S2T copy of SFA/SFB
#
tiled_copy_s2t_sfa, tCsSFA_compact_s2t, tCtSFA_compact_s2t = (
self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA)
)
tiled_copy_s2t_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = (
self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB)
)
#
# Persistent tile scheduling loop
#
tile_sched = utils.StaticPersistentTileScheduler.create(
tile_sched_params, cute.arch.block_idx(), grid_dim
)
# grouped gemm tile scheduler helper will compute the group index for the tile we're working on
group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper(
group_count,
tile_sched_params,
self.cluster_tile_shape_mnk,
utils.create_initial_search_state(),
)
work_tile = tile_sched.initial_work_tile_info()
ab_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_ab_stage
)
acc_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_acc_stage
)
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_tile_cnt,
cur_group_idx,
) = group_gemm_ts_helper.search_cluster_tile_count_k(
cur_tile_coord,
problem_sizes_mnkl,
)
# (MMA, MMA_M, MMA_N)
tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)]
# Peek (try_wait) AB buffer full for k_tile = 0
ab_consumer_state.reset_count()
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < cur_k_tile_cnt and is_leader_cta:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
)
#
# Wait for accumulator buffer empty
#
if is_leader_cta:
acc_pipeline.producer_acquire(acc_producer_state)
#
# Reset the ACCUMULATE field for each tile
#
tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
#
# Mma mainloop
#
for k_tile in range(cur_k_tile_cnt):
if is_leader_cta:
# Conditionally wait for AB buffer full
ab_pipeline.consumer_wait(
ab_consumer_state, peek_ab_full_status
)
# Copy SFA/SFB from smem to tmem
s2t_stage_coord = (
None,
None,
None,
None,
ab_consumer_state.index,
)
tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
cute.copy(
tiled_copy_s2t_sfa,
tCsSFA_compact_s2t_staged,
tCtSFA_compact_s2t,
)
cute.copy(
tiled_copy_s2t_sfb,
tCsSFB_compact_s2t_staged,
tCtSFB_compact_s2t,
)
# tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB
num_kblocks = cute.size(tCrA, mode=[2])
for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
kblock_coord = (
None,
None,
kblock_idx,
ab_consumer_state.index,
)
# Set SFA/SFB tensor to tiled_mma
sf_kblock_coord = (None, None, kblock_idx)
tiled_mma.set(
tcgen05.Field.SFA,
tCtSFA[sf_kblock_coord].iterator,
)
tiled_mma.set(
tcgen05.Field.SFB,
tCtSFB[sf_kblock_coord].iterator,
)
cute.gemm(
tiled_mma,
tCtAcc,
tCrA[kblock_coord],
tCrB[kblock_coord],
tCtAcc,
)
# Enable accumulate after the first kblock only.
if kblock_idx == 0:
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
# Async arrive AB buffer empty
ab_pipeline.consumer_release(ab_consumer_state)
# Peek (try_wait) AB buffer full for k_tile = k_tile + 1
ab_consumer_state.advance()
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < cur_k_tile_cnt:
if is_leader_cta:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
)
#
# Async arrive accumulator buffer full
#
if is_leader_cta:
acc_pipeline.producer_commit(acc_producer_state)
acc_producer_state.advance()
#
# Advance to next tile
#
tile_sched.advance_to_next_work()
work_tile = tile_sched.get_current_work()
#
# Wait for accumulator buffer empty
#
acc_pipeline.producer_tail(acc_producer_state)
#
# Specialized epilogue warps
#
if warp_idx < self.mma_warp_id:
# initialize tensorap for C
tensormap_manager.init_tensormap_from_atom(
tma_atom_c,
tensormap_c_smem_ptr,
self.epilog_warp_id[0],
)
#
# Alloc tensor memory buffer
#
if warp_idx == self.epilog_warp_id[0]:
cute.arch.alloc_tmem(
num_tmem_alloc_cols,
tmem_holding_buf,
is_two_cta=use_2cta_instrs,
)
#
# Bar sync for retrieve tensor memory ptr from shared memory
#
self.tmem_alloc_barrier.arrive_and_wait()
#
# Retrieving tensor memory ptr and make accumulator tensor
#
acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
self.acc_dtype,
alignment=16,
ptr_to_buffer_holding_addr=tmem_holding_buf,
)
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
### Start from here
#
# 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
)
)
tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, c_dtype)
tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(
tiled_copy_t2r, tTR_rC, epi_tidx, sC
)
tma_atom_c, bSG_sC, bSG_gC_partitioned = (
self.epilog_gmem_copy_and_partition(
epi_tidx, tma_atom_c, tCgC, epi_tile, sC
)
)
#
# Persistent tile scheduling loop
#
tile_sched = utils.StaticPersistentTileScheduler.create(
tile_sched_params, cute.arch.block_idx(), grid_dim
)
# grouped gemm tile scheduler helper will compute the group index for the tile we're working on
group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper(
group_count,
tile_sched_params,
self.cluster_tile_shape_mnk,
utils.create_initial_search_state(),
)
work_tile = tile_sched.initial_work_tile_info()
acc_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
)
# Threads/warps participating in tma store pipeline
c_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread,
32 * len(self.epilog_warp_id),
)
c_pipeline = pipeline.PipelineTmaStore.create(
num_stages=self.num_c_stage,
producer_group=c_producer_group,
)
# group index to start searching
last_group_idx = cutlass.Int32(-1)
while work_tile.is_valid_tile:
cur_tile_coord = work_tile.tile_idx
grouped_gemm_cta_tile_info = group_gemm_ts_helper.delinearize_z(
cur_tile_coord,
problem_sizes_mnkl,
)
cur_group_idx = grouped_gemm_cta_tile_info.group_idx
is_group_changed = cur_group_idx != last_group_idx
if is_group_changed:
# construct tensor c based on real shape, stride information
real_tensor_c = self.make_tensor_abc_for_tensormap_update(
cur_group_idx,
c_dtype,
(
grouped_gemm_cta_tile_info.problem_shape_m,
grouped_gemm_cta_tile_info.problem_shape_n,
grouped_gemm_cta_tile_info.problem_shape_k,
),
strides_abc,
ptrs_abc,
2, # 2 for tensor C
)
tensormap_manager.update_tensormap(
((real_tensor_c),),
((tma_atom_c),),
((tensormap_c_gmem_ptr),),
self.epilog_warp_id[0],
(tensormap_c_smem_ptr,),
)
mma_tile_coord_mnl = (
grouped_gemm_cta_tile_info.cta_tile_idx_m
// cute.size(tiled_mma.thr_id.shape),
grouped_gemm_cta_tile_info.cta_tile_idx_n,
0,
)
cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k
#
# Slice to per mma tile index
#
# ((ATOM_V, REST_V), EPI_M, EPI_N)
bSG_gC = bSG_gC_partitioned[
(
None,
None,
None,
*mma_tile_coord_mnl,
)
]
# 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)
]
#
# Wait for accumulator buffer full
#
acc_pipeline.consumer_wait(acc_consumer_state)
tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
if is_group_changed:
if warp_idx == self.epilog_warp_id[0]:
tensormap_manager.fence_tensormap_update(tensormap_c_gmem_ptr)
#
# Store accumulator to global memory in subtiles
#
subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt
tma_desc_ptr_c = tensormap_manager.get_tensormap_ptr(
tensormap_c_gmem_ptr,
cute.AddressSpace.generic,
)
tRS_rAcc = tiled_copy_r2s.retile(tTR_rAcc)
for subtile_idx in range(subtile_cnt):
#
# Load accumulator from tensor memory buffer to register
#
tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
#
# Convert to C type
#
acc_vec = tRS_rAcc.load()
tRS_rC.store(acc_vec.to(c_dtype))
#
# Store C to shared memory
#
c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage
cute.copy(
tiled_copy_r2s,
tRS_rC,
tRS_sC[(None, None, None, c_buffer)],
)
# Fence and barrier to make sure shared memory store is visible to TMA store
cute.arch.fence_proxy(
cute.arch.ProxyKind.async_shared,
space=cute.arch.SharedSpace.shared_cta,
)
self.epilog_sync_barrier.arrive_and_wait()
#
# TMA store C to global memory
#
if warp_idx == self.epilog_warp_id[0]:
cute.copy(
tma_atom_c,
bSG_sC[(None, c_buffer)],
bSG_gC[(None, subtile_idx)],
tma_desc_ptr=tma_desc_ptr_c,
)
# Fence and barrier to make sure shared memory store is visible to TMA store
c_pipeline.producer_commit()
c_pipeline.producer_acquire()
self.epilog_sync_barrier.arrive_and_wait()
#
# Async arrive accumulator buffer empty
#
with cute.arch.elect_one():
acc_pipeline.consumer_release(acc_consumer_state)
acc_consumer_state.advance()
#
# Advance to next tile
#
tile_sched.advance_to_next_work()
work_tile = tile_sched.get_current_work()
last_group_idx = cur_group_idx
#
# Dealloc the tensor memory buffer
#
if warp_idx == self.epilog_warp_id[0]:
cute.arch.relinquish_tmem_alloc_permit(is_two_cta=use_2cta_instrs)
self.epilog_sync_barrier.arrive_and_wait()
if warp_idx == self.epilog_warp_id[0]:
if use_2cta_instrs:
cute.arch.mbarrier_arrive(
tmem_dealloc_mbar_ptr, cta_rank_in_cluster ^ 1
)
cute.arch.mbarrier_wait(tmem_dealloc_mbar_ptr, 0)
cute.arch.dealloc_tmem(
acc_tmem_ptr, num_tmem_alloc_cols, is_two_cta=use_2cta_instrs
)
#
# Wait for C store complete
#
c_pipeline.producer_tail()
@cute.jit
def __call__(
self,
ptr_of_tensor_of_problem_sizes: cute.Pointer,
ptr_of_tensor_of_abc_ptrs: cute.Pointer,
ptr_of_tensor_of_sfasfb_ptrs: cute.Pointer,
ptr_of_tensor_of_tensormap: cute.Pointer,
ptr_of_tensor_strides_abc: cute.Pointer,
sm_count: cutlass.Constexpr[int],
total_num_clusters: cutlass.Constexpr[int],
num_groups: cutlass.Constexpr[int],
max_active_clusters: cutlass.Constexpr[int],
):
tensor_of_abc_ptrs = cute.make_tensor(
ptr_of_tensor_of_abc_ptrs, cute.make_layout((num_groups, 3), stride=(3, 1))
)
tensor_of_sfasfb_ptrs = cute.make_tensor(
ptr_of_tensor_of_sfasfb_ptrs, cute.make_layout((num_groups, 2), stride=(2, 1))
)
problem_shape_mnkl = cute.make_tensor(
ptr_of_tensor_of_problem_sizes, cute.make_layout((num_groups, 4), stride=(4, 1))
)
tensormap_stride1 = Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8
tensormap_stride0 = Sm100GroupedBlockScaledGemmKernel.num_tensormaps * tensormap_stride1
tensor_of_tensormap = cute.make_tensor(
ptr_of_tensor_of_tensormap,
cute.make_layout(
(sm_count, Sm100GroupedBlockScaledGemmKernel.num_tensormaps, tensormap_stride1),
stride=(tensormap_stride0, tensormap_stride1, 1),
),
)
# provides 3D logical view onto 1D memory
# - The pointer gives the base address of the flat data in memory.
# - The layout (shape + stride) tells CuTe how to translate a logical 3‑D index into a linear offset from that base.
tensor_of_strides_abc = cute.make_tensor(
ptr_of_tensor_strides_abc, cute.make_layout((num_groups, 3, 2), stride=(6, 2, 1))
)
# Use fake shape for initial Tma descriptor and atom setup
# The real Tma desc and atom will be updated during kernel execution.
min_a_shape = (cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(1))
min_b_shape = (cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(1))
min_c_shape = (cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(1))
initial_a = cute.make_tensor(
cute.make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16,),
cute.make_layout(
(min_a_shape[0], cute.assume(min_a_shape[2], 32), min_a_shape[3]),
stride=(
cute.assume(min_a_shape[2], 32),
1,
cute.assume(min_a_shape[0] * min_a_shape[2], 32),
),
),
)
initial_b = cute.make_tensor(
cute.make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16,),
cute.make_layout(
(min_b_shape[1], cute.assume(min_b_shape[2], 32), min_b_shape[3]),
stride=(
cute.assume(min_b_shape[2], 32),
1,
cute.assume(min_b_shape[1] * min_b_shape[2], 32),
),
),
)
initial_c = cute.make_tensor(
cute.make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16,),
cute.make_layout(
(min_c_shape[1], cute.assume(min_c_shape[2], 32), min_c_shape[3]),
stride=(
cute.assume(min_c_shape[2], 32),
1,
cute.assume(min_c_shape[1] * min_c_shape[2], 32),
),
),
)
# Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout
# ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL)
sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
initial_a.shape, sf_vec_size
)
# ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
initial_b.shape, sf_vec_size
)
# Create initial SFA and SFB tensors with fake shape and null pointer.
initial_sfa = cute.make_tensor(
cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16,), sfa_layout)
initial_sfb = cute.make_tensor(
cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16,), sfb_layout)
a_major_mode = utils.LayoutEnum.from_tensor(initial_a).mma_major_mode()
b_major_mode = utils.LayoutEnum.from_tensor(initial_b).mma_major_mode()
self.c_layout = utils.LayoutEnum.from_tensor(initial_c)
mma_inst_shape_mn = (
self.mma_tiler[0],
self.mma_tiler[1],
)
mma_inst_shape_mn_sfb = (
mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1),
cute.round_up(mma_inst_shape_mn[1], 128),
)
tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
ab_dtype,
a_major_mode,
b_major_mode,
sf_dtype,
sf_vec_size,
self.cta_group,
mma_inst_shape_mn,
)
tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
ab_dtype,
a_major_mode,
b_major_mode,
sf_dtype,
sf_vec_size,
cute.nvgpu.tcgen05.CtaGroup.ONE,
mma_inst_shape_mn_sfb,
)
atom_thr_size = cute.size(tiled_mma.thr_id.shape)
# Compute mma instruction shapes
# (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K)
self.mma_inst_shape_mn = (
self.mma_tiler[0],
self.mma_tiler[1],
)
# (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K)
self.mma_inst_shape_mn_sfb = (
self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1),
cute.round_up(self.mma_inst_shape_mn[1], 128),
)
self.mma_tiler = (
self.mma_inst_shape_mn[0],
self.mma_inst_shape_mn[1],
mma_inst_shape_k * mma_inst_tile_k,
)
self.mma_tiler_sfb = (
self.mma_inst_shape_mn_sfb[0],
self.mma_inst_shape_mn_sfb[1],
mma_inst_shape_k * mma_inst_tile_k,
)
self.cta_tile_shape_mnk = (
self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
self.mma_tiler[1],
self.mma_tiler[2],
)
self.cluster_tile_shape_mnk = tuple(
x * y for x, y in zip(self.cta_tile_shape_mnk, (*self.cluster_shape_mn, 1))
)
# **Why* we compute the shape and what it's used for*
# Because the epilogue needs a concrete tile shape to safely and efficiently write results back to global memory.
# The MMA tile size and the optimal store tile size aren’t always the same—store needs to respect C’s layout, alignment, and vectorization/TMA requirements.
# The “epilogue tile” tells the kernel how to break the accumulator into chunks, how big sC should be in shared memory,
# and how the TMA store should be issued without out‑of‑bounds or misaligned accesses.
epi_tile = sm100_utils.compute_epilogue_tile_shape(
self.cta_tile_shape_mnk,
self.use_2cta_instrs,
self.c_layout,
c_dtype,
)
# Stage counts are computed to maximize pipelining without exceeding SMEM: we estimate SMEM per A/B/SF stage and per C stage,
# then choose the largest num_ab_stage that fits (given occupancy and reserved bytes), and use any leftover SMEM to add more C stages.
# This balances latency hiding with SMEM limits and avoids over-allocation.
self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages(
tiled_mma,
self.mma_tiler,
ab_dtype,
ab_dtype,
epi_tile,
c_dtype,
self.c_layout,
sf_dtype,
sf_vec_size,
smem_capacity,
self.occupancy,
)
# Compute A/B/SFA/SFB/C shared memory layout
# Because the SMEM layout for A is staged—it includes an extra “stage” dimension for pipelining (double/triple buffering).
# The number of stages tells the layout how many A tiles to allocate in shared memory and how to index them.
# So num_ab_stage directly determines the SMEM footprint and addressing for A (and similarly B/SF)
self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
tiled_mma,
self.mma_tiler,
ab_dtype,
self.num_ab_stage,
)
self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
tiled_mma,
self.mma_tiler,
ab_dtype,
self.num_ab_stage,
)
self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
tiled_mma,
self.mma_tiler,
sf_vec_size,
self.num_ab_stage,
)
self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
tiled_mma,
self.mma_tiler,
sf_vec_size,
self.num_ab_stage,
)
self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(
c_dtype,
self.c_layout,
epi_tile,
self.num_c_stage,
)
# Define shared storage for kernel
@cute.struct
class SharedStorage:
# cute.struct.MemRange[T, N] is a typed shared‑memory buffer of N elements of type T, not bytes.
# The size N is a count of elements. Total bytes = N * sizeof(T).
tensormap_buffer: cute.struct.MemRange[
cutlass.Int64, self.size_tensormap_in_i64
]
ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
# tmem_dealloc_mbar_ptr is the shared‑memory barrier (mbarrier) used to synchronize TMEM deallocation across warps/CTAs.
# It ensures everyone is done using tensor memory before it’s freed.
# It’s a single 64‑bit slot because mbarriers are represented as 64‑bit values in SMEM.
tmem_dealloc_mbar_ptr: cutlass.Int64
# tmem_holding_buf is a small shared‑memory slot used to hold the pointer to tensor memory (TMEM) that’s allocated at runtime.
tmem_holding_buf: cutlass.Int32
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sC: cute.struct.Align[
cute.struct.MemRange[
c_dtype,
cute.cosize(self.c_smem_layout_staged.outer),
],
buffer_align_bytes,
]
# (MMA, MMA_M, MMA_K, STAGE)
sA: cute.struct.Align[
cute.struct.MemRange[
ab_dtype, cute.cosize(self.a_smem_layout_staged.outer)
],
buffer_align_bytes,
]
# (MMA, MMA_N, MMA_K, STAGE)
sB: cute.struct.Align[
cute.struct.MemRange[
ab_dtype, cute.cosize(self.b_smem_layout_staged.outer)
],
buffer_align_bytes,
]
# (MMA, MMA_M, MMA_K, STAGE)
sSFA: cute.struct.Align[
cute.struct.MemRange[
sf_dtype, cute.cosize(self.sfa_smem_layout_staged)
],
buffer_align_bytes,
]
# (MMA, MMA_N, MMA_K, STAGE)
sSFB: cute.struct.Align[
cute.struct.MemRange[
sf_dtype, cute.cosize(self.sfb_smem_layout_staged)
],
buffer_align_bytes,
]
self.shared_storage = SharedStorage
# Compute cluster layout
self.cluster_layout_vmnk = cute.tiled_divide(
cute.make_layout((*self.cluster_shape_mn, 1)),
(tiled_mma.thr_id.shape,),
)
self.cluster_layout_sfb_vmnk = cute.tiled_divide(
cute.make_layout((*self.cluster_shape_mn, 1)),
(tiled_mma_sfb.thr_id.shape,),
)
# Compute number of multicast CTAs for A/B
self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1])
self.is_a_mcast = self.num_mcast_ctas_a > 1
self.is_b_mcast = self.num_mcast_ctas_b > 1
self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1
# Setup TMA load for A
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_tiled_tma_atom_A(
a_op,
initial_a,
a_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
)
# Setup TMA load for B
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_tiled_tma_atom_B(
b_op,
initial_b,
b_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
)
# Setup TMA load for SFA
sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(
self.cluster_shape_mn, tiled_mma.thr_id
)
sfa_smem_layout = cute.slice_(
self.sfa_smem_layout_staged, (None, None, None, 0)
)
tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
sfa_op,
initial_sfa,
sfa_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
internal_type=cutlass.Int16,
)
# Setup TMA load for SFB
sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(
self.cluster_shape_mn, tiled_mma.thr_id
)
sfb_smem_layout = cute.slice_(
self.sfb_smem_layout_staged, (None, None, None, 0)
)
tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
sfb_op,
initial_sfb,
sfb_smem_layout,
self.mma_tiler_sfb,
tiled_mma_sfb,
self.cluster_layout_sfb_vmnk.shape,
internal_type=cutlass.Int16,
)
a_copy_size = cute.size_in_bytes(ab_dtype, a_smem_layout)
b_copy_size = cute.size_in_bytes(ab_dtype, b_smem_layout)
sfa_copy_size = cute.size_in_bytes(sf_dtype, sfa_smem_layout)
sfb_copy_size = cute.size_in_bytes(sf_dtype, sfb_smem_layout)
self.num_tma_load_bytes = (
a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size
) * atom_thr_size
# Setup TMA store for C
epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0))
tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
cpasync.CopyBulkTensorTileS2GOp(),
initial_c,
epi_smem_layout,
epi_tile,
)
# Compute grid size (For first problem)
# total_num_clusters = 352
# self.cluster_shape_mn = (1, 2)
# max_active_clusters = 148
self.tile_sched_params, grid = self._compute_grid(
total_num_clusters, self.cluster_shape_mn, max_active_clusters
)
# cute.printf("launch kernel\n")
# cute.printf("grid: {}\n", grid),
# cute.printf("block: [{}, 1, 1]\n", self.threads_per_cta)
# cute.printf("cluster: {}\n", self.cluster_shape_mn)
# cute.printf("num ab stage: {}\n", self.num_ab_stage)
# cute.printf("num acc stage: {}\n", self.num_acc_stage)
# cute.printf("num c stage: {}\n", self.num_c_stage)
# Launch the kernel synchronously
self.ported_kernel(
tiled_mma=tiled_mma,
tiled_mma_sfb=tiled_mma_sfb,
tma_atom_a=tma_atom_a,
mA_mkl=tma_tensor_a,
tma_atom_b=tma_atom_b,
mB_nkl=tma_tensor_b,
tma_atom_sfa=tma_atom_sfa,
mSFA_mkl=tma_tensor_sfa,
tma_atom_sfb=tma_atom_sfb,
mSFB_nkl=tma_tensor_sfb,
tma_atom_c=tma_atom_c,
mC_mnl=tma_tensor_c,
cluster_layout_vmnk=self.cluster_layout_vmnk,
cluster_layout_sfb_vmnk=self.cluster_layout_sfb_vmnk,
a_smem_layout_staged=self.a_smem_layout_staged,
b_smem_layout_staged=self.b_smem_layout_staged,
sfa_smem_layout_staged=self.sfa_smem_layout_staged,
sfb_smem_layout_staged=self.sfb_smem_layout_staged,
c_smem_layout_staged=self.c_smem_layout_staged,
epi_tile=epi_tile,
tile_sched_params=self.tile_sched_params,
group_count=num_groups,
problem_sizes_mnkl=problem_shape_mnkl,
ptrs_abc=tensor_of_abc_ptrs,
ptrs_sfasfb=tensor_of_sfasfb_ptrs,
tensormaps=tensor_of_tensormap,
strides_abc=tensor_of_strides_abc
# might be just this So if you choose A=MKL, B=NKL, C=MNL (i.e., A is M‑major, B is N‑major, C is M‑major), then yes—is_mode0_major is True for all three, and the stride for each group becomes (1, mode0).
# strides_abc,
# tensor_address_abc,
# tensor_address_sfasfb,
# tensormap_cute_tensor,
).launch(
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
smem=self.shared_storage.size_in_bytes(),
min_blocks_per_mp=1,
)
return
@cute.jit
def make_tensor_abc_for_tensormap_update(
self,
group_idx: cutlass.Int32,
dtype: Type[cutlass.Numeric],
problem_shape_mnk: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32],
strides_abc: cute.Tensor,
tensor_address_abc: cute.Tensor,
tensor_index: int,
):
"""Extract stride and tensor address for a given group and construct a global tensor for A, B or C.
This function is used within the kernel to dynamically create a CUTE tensor
representing A, B, or C for the current group being processed, using the
group-specific address, shape, and stride information.
:param group_idx: The index of the current group within the grouped GEMM.
:type group_idx: cutlass.Int32
:param dtype: The data type of the tensor elements (e.g., cutlass.Float16).
:type dtype: Type[cutlass.Numeric]
:param problem_shape_mnk: The (M, N, K) problem shape for the current group.
:type problem_shape_mnk: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32]
:param strides_abc: Tensor containing strides for A, B, C for all groups. Layout: (group_count, 3, 2).
:type strides_abc: cute.Tensor
:param tensor_address_abc: Tensor containing global memory addresses for A, B, C for all groups. Layout: (group_count, 3).
:type tensor_address_abc: cute.Tensor
:param tensor_index: Specifies which tensor to create: 0 for A, 1 for B, 2 for C.
:type tensor_index: int
:return: A CUTE tensor representing the requested global memory tensor (A, B, or C) for the specified group.
:rtype: cute.Tensor
:raises TypeError: If the provided dtype is not a subclass of cutlass.Numeric.
"""
ptr_i64 = tensor_address_abc[(group_idx, tensor_index)]
if cutlass.const_expr(
not isclass(dtype) or not issubclass(dtype, cutlass.Numeric)
):
raise TypeError(
f"dtype must be a type of cutlass.Numeric, got {type(dtype)}"
)
tensor_gmem_ptr = cute.make_ptr(
dtype, ptr_i64, cute.AddressSpace.gmem, assumed_align=16
)
strides_tensor_gmem = strides_abc[(group_idx, tensor_index, None)]
strides_tensor_reg = cute.make_rmem_tensor(
cute.make_layout(2),
strides_abc.element_type,
)
cute.autovec_copy(strides_tensor_gmem, strides_tensor_reg)
stride_mn = strides_tensor_reg[0]
stride_k = strides_tensor_reg[1]
c1 = cutlass.Int32(1)
c0 = cutlass.Int32(0)
if cutlass.const_expr(tensor_index == 0): # tensor A
m = problem_shape_mnk[0]
k = problem_shape_mnk[2]
return cute.make_tensor(
tensor_gmem_ptr,
cute.make_layout((m, k, c1), stride=(stride_mn, stride_k, c0)),
)
elif cutlass.const_expr(tensor_index == 1): # tensor B
n = problem_shape_mnk[1]
k = problem_shape_mnk[2]
return cute.make_tensor(
tensor_gmem_ptr,
cute.make_layout((n, k, c1), stride=(stride_mn, stride_k, c0)),
)
else: # tensor C
m = problem_shape_mnk[0]
n = problem_shape_mnk[1]
return cute.make_tensor(
tensor_gmem_ptr,
cute.make_layout((m, n, c1), stride=(stride_mn, stride_k, c0)),
)
@cute.jit
def make_tensor_sfasfb_for_tensormap_update(
self,
group_idx: cutlass.Int32,
dtype: Type[cutlass.Numeric],
problem_shape_mnk: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32],
tensor_address_sfasfb: cute.Tensor,
tensor_index: int,
):
"""Extract tensor address for a given group and construct a global tensor for SFA or SFB.
This function is used within the kernel to dynamically create a CUTE tensor
representing SFA or SFB for the current group being processed, using the
group-specific address, shape information.
:param group_idx: The index of the current group within the grouped GEMM.
:type group_idx: cutlass.Int32
:param dtype: The data type of the tensor elements (e.g., cutlass.Float16).
:type dtype: Type[cutlass.Numeric]
:param problem_shape_mnk: The (M, N, K) problem shape for the current group.
:type problem_shape_mnk: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32]
:param tensor_address_sfasfb: Tensor containing global memory addresses for SFA, SFB for all groups. Layout: (group_count, 2).
:type tensor_address_sfasfb: cute.Tensor
:param tensor_index: Specifies which tensor to create: 0 for SFA, 1 for SFB.
:type tensor_index: int
:return: A CUTE tensor representing the requested global memory tensor (SFA, SFB) for the specified group.
:rtype: cute.Tensor
:raises TypeError: If the provided dtype is not a subclass of cutlass.Numeric.
"""
ptr_i64 = tensor_address_sfasfb[(group_idx, tensor_index)]
if cutlass.const_expr(
not isclass(dtype) or not issubclass(dtype, cutlass.Numeric)
):
raise TypeError(
f"dtype must be a type of cutlass.Numeric, got {type(dtype)}"
)
tensor_gmem_ptr = cute.make_ptr(
dtype, ptr_i64, cute.AddressSpace.gmem, assumed_align=16
)
c1 = cutlass.Int32(1)
if cutlass.const_expr(tensor_index == 0): # tensor SFA
m = problem_shape_mnk[0]
k = problem_shape_mnk[2]
sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
(m, k, c1), sf_vec_size
)
return cute.make_tensor(
tensor_gmem_ptr,
sfa_layout,
)
else: # tensor SFB
n = problem_shape_mnk[1]
k = problem_shape_mnk[2]
sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
(n, k, c1), sf_vec_size
)
return cute.make_tensor(
tensor_gmem_ptr,
sfb_layout,
)
@staticmethod
def _compute_stages(
tiled_mma: cute.TiledMma,
mma_tiler_mnk: Tuple[int, int, int],
a_dtype: Type[cutlass.Numeric],
b_dtype: Type[cutlass.Numeric],
epi_tile: cute.Tile,
c_dtype: Type[cutlass.Numeric],
c_layout: utils.LayoutEnum,
sf_dtype: Type[cutlass.Numeric],
sf_vec_size: int,
smem_capacity: int,
occupancy: int,
) -> Tuple[int, int, int]:
"""Computes the number of stages for A/B/C operands based on heuristics.
: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 epi_tile: The epilogue tile shape.
:type epi_tile: cute.Tile
:param c_dtype: Data type of operand C (output).
:type c_dtype: type[cutlass.Numeric]
:param c_layout: Layout enum of operand C.
:type c_layout: utils.LayoutEnum
:param sf_dtype: Data type of Scale factor.
:type sf_dtype: type[cutlass.Numeric]
:param sf_vec_size: Scale factor vector size.
:type sf_vec_size: 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
:return: A tuple containing the computed number of stages for:
(ACC stages, A/B operand stages, C stages)
:rtype: tuple[int, int, int]
"""
# ACC stages
num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2
# Default C stages
num_c_stage = 2
# Calculate smem layout and size for one stage of A, B, SFA, SFB and C
a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(
tiled_mma,
mma_tiler_mnk,
a_dtype,
1, # a tmp 1 stage is provided
)
b_smem_layout_staged_one = sm100_utils.make_smem_layout_b(
tiled_mma,
mma_tiler_mnk,
b_dtype,
1, # a tmp 1 stage is provided
)
sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa(
tiled_mma,
mma_tiler_mnk,
sf_vec_size,
1, # a tmp 1 stage is provided
)
sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb(
tiled_mma,
mma_tiler_mnk,
sf_vec_size,
1, # a tmp 1 stage is provided
)
c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi(
c_dtype,
c_layout,
epi_tile,
1,
)
ab_bytes_per_stage = (
cute.size_in_bytes(a_dtype, a_smem_layout_stage_one)
+ cute.size_in_bytes(b_dtype, b_smem_layout_staged_one)
+ cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one)
+ cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one)
)
mbar_helpers_bytes = 1024
c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one)
c_bytes = c_bytes_per_stage * num_c_stage
# Calculate A/B/SFA/SFB stages:
# Start with total smem per CTA (capacity / occupancy)
# Subtract reserved bytes and initial C stages bytes
# Divide remaining by bytes needed per A/B/SFA/SFB stage
num_ab_stage = (
smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes)
) // ab_bytes_per_stage
# Refine epilogue stages:
# Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes
# Add remaining unused smem to epilogue
num_c_stage += (
smem_capacity
- occupancy * ab_bytes_per_stage * num_ab_stage
- occupancy * (mbar_helpers_bytes + c_bytes)
) // (occupancy * c_bytes_per_stage)
return num_acc_stage, num_ab_stage, num_c_stage
@staticmethod
def _compute_grid(
total_num_clusters: int,
cluster_shape_mn: tuple[int, int],
max_active_clusters: cutlass.Constexpr[int],
) -> tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]:
"""Compute tile scheduler parameters and grid shape for grouped GEMM operations.
:param total_num_clusters: Total number of clusters to process across all groups.
:type total_num_clusters: int
:param cluster_shape_mn: Shape of each cluster in M, N dimensions.
:type cluster_shape_mn: tuple[int, int]
:param max_active_clusters: Maximum number of active clusters.
:type max_active_clusters: cutlass.Constexpr[int]
:return: A tuple containing:
- tile_sched_params: Parameters for the persistent tile scheduler.
- grid: Grid shape for kernel launch.
:rtype: tuple[utils.PersistentTileSchedulerParams, tuple[int, ...]]
"""
# Create problem shape with M, N dimensions from cluster shape
# and L dimension representing the total number of clusters.
problem_shape_ntile_mnl = (
cluster_shape_mn[0],
cluster_shape_mn[1],
cutlass.Int32(total_num_clusters),
)
# (128, 128, 352)
tile_sched_params = utils.PersistentTileSchedulerParams(
problem_shape_ntile_mnl, (*cluster_shape_mn, 1)
)
grid = utils.StaticPersistentTileScheduler.get_grid_shape(
tile_sched_params, max_active_clusters
)
return tile_sched_params, grid
@staticmethod
def _get_mbar_smem_bytes(**kwargs_stages: int) -> int:
"""Calculate shared memory consumption for memory barriers based on provided stages.
Each stage requires 2 barriers, and each barrier consumes 8 bytes of shared memory.
The total consumption is the sum across all provided stages. This function calculates the total
shared memory needed for these barriers.
:param kwargs_stages: Variable keyword arguments where each key is a stage name
(e.g., num_acc_stage, num_ab_stage) and each value is the
number of stages of that type.
:type kwargs_stages: int
:return: Total shared memory bytes required for all memory barriers.
:rtype: int
"""
num_barriers_per_stage = 2
num_bytes_per_barrier = 8
mbar_smem_consumption = sum(
[
num_barriers_per_stage * num_bytes_per_barrier * stage
for stage in kwargs_stages.values()
]
)
return mbar_smem_consumption
def mainloop_s2t_copy_and_partition(
self,
sSF: cute.Tensor,
tSF: cute.Tensor,
) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
"""
Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination).
:param sSF: The scale factor tensor in smem
:type sSF: cute.Tensor
:param tSF: The scale factor tensor in tmem
:type tSF: cute.Tensor
:return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where:
- tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t)
- tCsSF_compact_s2t: The partitioned scale factor tensor in smem
- tSF_compact_s2t: The partitioned scale factor tensor in tmem
:rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
"""
# (MMA, MMA_MN, MMA_K, STAGE)
tCsSF_compact = cute.filter_zeros(sSF)
# (MMA, MMA_MN, MMA_K)
tCtSF_compact = cute.filter_zeros(tSF)
# Make S2T CopyAtom and tiledCopy
copy_atom_s2t = cute.make_copy_atom(
tcgen05.Cp4x32x128bOp(self.cta_group),
sf_dtype,
)
tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact)
thr_copy_s2t = tiled_copy_s2t.get_slice(0)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
tiled_copy_s2t, tCsSF_compact_s2t_
)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact)
return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t
def epilog_tmem_copy_and_partition(
self,
tidx: cutlass.Int32,
tAcc: cute.Tensor,
gC_mnl: cute.Tensor,
epi_tile: cute.Tile,
use_2cta_instrs: Union[cutlass.Boolean, bool],
) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
"""
Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination).
:param tidx: The thread index in epilogue warp groups
:type tidx: cutlass.Int32
:param tAcc: The accumulator tensor to be copied and partitioned
:type tAcc: cute.Tensor
:param gC_mnl: The global tensor C
:type gC_mnl: cute.Tensor
:param epi_tile: The epilogue tiler
:type epi_tile: cute.Tile
:param use_2cta_instrs: Whether use_2cta_instrs is enabled
:type use_2cta_instrs: bool
:return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where:
- tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r)
- tTR_tAcc: The partitioned accumulator tensor
- tTR_rAcc: The accumulated tensor in register used to hold t2r results
:rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
"""
# Make tiledCopy for tensor memory load
copy_atom_t2r = sm100_utils.get_tmem_load_op(
self.cta_tile_shape_mnk,
self.c_layout,
c_dtype,
self.acc_dtype,
epi_tile,
use_2cta_instrs,
)
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE)
tAcc_epi = cute.flat_divide(
tAcc[((None, None), 0, 0, None)],
epi_tile,
)
# (EPI_TILE_M, EPI_TILE_N)
tiled_copy_t2r = tcgen05.make_tmem_copy(
copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]
)
thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
# (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, 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, RestM, RestN, RestL)
tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
# (T2R, T2R_M, T2R_N)
tTR_rAcc = cute.make_rmem_tensor(
tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype
)
return tiled_copy_t2r, tTR_tAcc, tTR_rAcc
def epilog_smem_copy_and_partition(
self,
tiled_copy_t2r: cute.TiledCopy,
tTR_rC: cute.Tensor,
tidx: cutlass.Int32,
sC: cute.Tensor,
) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
"""
Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination).
:param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r)
:type tiled_copy_t2r: cute.TiledCopy
:param tTR_rC: The partitioned accumulator tensor
:type tTR_rC: cute.Tensor
:param tidx: The thread index in epilogue warp groups
:type tidx: cutlass.Int32
:param sC: The shared memory tensor to be copied and partitioned
:type sC: cute.Tensor
:type sepi: cute.Tensor
:return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where:
- tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s)
- tRS_rC: The partitioned tensor C (register source)
- 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(
self.c_layout, c_dtype, self.acc_dtype, tiled_copy_t2r
)
tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
# (R2S, R2S_M, R2S_N, PIPE_D)
thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
tRS_sC = thr_copy_r2s.partition_D(sC)
# (R2S, R2S_M, R2S_N)
tRS_rC = tiled_copy_r2s.retile(tTR_rC)
return tiled_copy_r2s, tRS_rC, tRS_sC
def epilog_gmem_copy_and_partition(
self,
tidx: cutlass.Int32,
atom: Union[cute.CopyAtom, cute.TiledCopy],
gC_mnl: cute.Tensor,
epi_tile: cute.Tile,
sC: cute.Tensor,
) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]:
"""Make tiledCopy for global memory store, then use it to:
partition shared memory (source) and global memory (destination) for TMA store version.
:param tidx: The thread index in epilogue warp groups
:type tidx: cutlass.Int32
:param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version
:type atom: cute.CopyAtom or cute.TiledCopy
:param gC_mnl: The global tensor C
:type gC_mnl: cute.Tensor
:param epi_tile: The epilogue tiler
:type epi_tile: cute.Tile
:param sC: The shared memory tensor to be copied and partitioned
:type sC: cute.Tensor
:return: A tuple containing (tma_atom_c, bSG_sC, bSG_gC) where:
- tma_atom_c: The TMA copy atom
- bSG_sC: The partitioned shared memory tensor C
- bSG_gC: The partitioned global tensor C
:rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]
"""
# (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
)
tma_atom_c = atom
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, RestM, RestN, RestL)
bSG_sC, bSG_gC = cpasync.tma_partition(
tma_atom_c,
0,
cute.make_layout(1),
sC_for_tma_partition,
gC_for_tma_partition,
)
return tma_atom_c, bSG_sC, bSG_gC
# Global caches to reduce repeated host-side setup overhead.
_compiled_kernel_cache = {}
_static_metadata_cache = {}
_tensormap_workspace_cache = {}
_pointer_metadata_cache = {}
_launch_context_cache = {}
_hardware_info = None
_sm_count_cache = None
_max_active_clusters_cache = {}
def make_problem_sizes_cache_key(problem_sizes: List[Tuple[int, int, int, int]]) -> str:
"""Build a deterministic cache key from the full ordered problem-size tuples."""
if not problem_sizes:
return "g0"
packed_groups = "|".join(f"{m},{n},{k},{l}" for m, n, k, l in problem_sizes)
return f"g{len(problem_sizes)}:{packed_groups}"
def _get_hardware_info() -> cutlass.utils.HardwareInfo:
global _hardware_info
if _hardware_info is None:
_hardware_info = cutlass.utils.HardwareInfo()
return _hardware_info
def _get_sm_count() -> int:
global _sm_count_cache
if _sm_count_cache is None:
_sm_count_cache = _get_hardware_info().get_max_active_clusters(1)
return _sm_count_cache
def _get_max_active_clusters(cluster_shape_mn: Tuple[int, int]) -> int:
cluster_size = cluster_shape_mn[0] * cluster_shape_mn[1]
if cluster_size not in _max_active_clusters_cache:
_max_active_clusters_cache[cluster_size] = (
_get_hardware_info().get_max_active_clusters(cluster_size)
)
return _max_active_clusters_cache[cluster_size]
def _get_static_metadata_tensors(
problem_sizes: List[Tuple[int, int, int, int]],
) -> Tuple[torch.Tensor, torch.Tensor]:
device_idx = torch.cuda.current_device()
cache_key = (device_idx, make_problem_sizes_cache_key(problem_sizes))
cached = _static_metadata_cache.get(cache_key)
if cached is not None:
return cached
# (num_groups, 4)
tensor_of_problem_sizes = torch.tensor(
problem_sizes, dtype=torch.int32, device="cuda"
)
# (num_groups, 3, 2)
strides_abc = [[(k, 1), (k, 1), (n, 1)] for _, n, k, _ in problem_sizes]
tensor_of_strides_abc = torch.tensor(strides_abc, dtype=torch.int32, device="cuda")
_static_metadata_cache[cache_key] = (
tensor_of_problem_sizes,
tensor_of_strides_abc,
)
return tensor_of_problem_sizes, tensor_of_strides_abc
def _get_pointer_metadata_tensors(
abc_ptrs,
sfasfb_ptrs,
) -> Tuple[torch.Tensor, torch.Tensor]:
device_idx = torch.cuda.current_device()
cache_key = (
device_idx,
tuple(abc_ptrs),
tuple(sfasfb_ptrs),
)
cached = _pointer_metadata_cache.get(cache_key)
if cached is not None:
return cached
tensor_of_abc_ptrs = torch.tensor(abc_ptrs, dtype=torch.int64, device="cuda")
tensor_of_sfasfb_ptrs = torch.tensor(sfasfb_ptrs, dtype=torch.int64, device="cuda")
_pointer_metadata_cache[cache_key] = (tensor_of_abc_ptrs, tensor_of_sfasfb_ptrs)
return tensor_of_abc_ptrs, tensor_of_sfasfb_ptrs
def _get_tensormap_workspace(sm_count: int) -> torch.Tensor:
device_idx = torch.cuda.current_device()
cache_key = (device_idx, sm_count)
tensormap_shape = (
sm_count,
Sm100GroupedBlockScaledGemmKernel.num_tensormaps,
Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8,
)
cached = _tensormap_workspace_cache.get(cache_key)
if cached is not None and tuple(cached.shape) == tensormap_shape:
return cached
tensor_of_tensormap = torch.empty(tensormap_shape, dtype=torch.int64, device="cuda")
_tensormap_workspace_cache[cache_key] = tensor_of_tensormap
return tensor_of_tensormap
# This function is used to compile the kernel once and cache it and then allow users to
# run the kernel multiple times to get more accurate timing results.
def compile_kernel(problem_sizes):
"""
Compile the kernel once and cache it using problem_sizes as the key.
This should be called before any timing measurements.
Returns:
The compiled kernel function
"""
global _compiled_kernel_cache
# Build cache key from the full ordered (m, n, k, l) tuples.
cache_key = make_problem_sizes_cache_key(problem_sizes)
# Check if we already have a compiled kernel for these problem sizes
if cache_key in _compiled_kernel_cache:
return _compiled_kernel_cache[cache_key]
# Dummy/null pointers are only used to define the kernel ABI for compile-time;
# real device pointers are passed at runtime from custom_kernel.
cute_ptr_of_tensor_of_problem_sizes = make_ptr(
cutlass.Int32, 0, cute.AddressSpace.gmem, assumed_align=16,
)
cute_ptr_of_tensor_of_abc_ptrs = make_ptr(
cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16,
)
cute_ptr_of_tensor_of_sfasfb_ptrs = make_ptr(
cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16,
)
# Each cluster needs its own set of tensormaps (one for A, B, SFA, SFB)
# Shape: (total_num_clusters, num_tensormaps=5, bytes_per_tensormap/8=16)
cute_ptr_of_tensor_of_tensormap = make_ptr(
cutlass.Int64, 0, cute.AddressSpace.gmem, assumed_align=16,
)
cute_ptr_of_tensor_of_strides_abc = make_ptr(
cutlass.Int32, 0, cute.AddressSpace.gmem, assumed_align=16,
)
kernel_runner = Sm100GroupedBlockScaledGemmKernel(
(128, 128),
cluster_shape_mn=cluster_shape_mn
)
# Fake cluster numbers for compile only.
# Compute the tile shape for each CUDA Thread Block (CTA)
# cta_tile_shape_mn: [M_tile, N_tile] = [128, 128] for this kernel
cta_tile_shape_mn = [128, mma_tiler_mnk[1]]
# cluster_tile_shape_mn: Total tile shape per cluster
cluster_tile_shape_mn = tuple(
x * y for x, y in zip(cta_tile_shape_mn, cluster_shape_mn)
)
# Compute total number of cluster tiles needed across all groups
# Each group's (m, n) dimensions are divided into tiles of size cluster_tile_shape_mn
# This determines the total grid size (bidz dimension) for kernel launch
total_num_clusters = 0
for m, n, _, _ in problem_sizes:
# Calculate number of tiles needed in M and N dimensions for this group
num_clusters_mn = tuple(
(x + y - 1) // y for x, y in zip((m, n), cluster_tile_shape_mn)
)
# Multiply M_tiles * N_tiles to get total tiles for this group
total_num_clusters += functools.reduce(lambda x, y: x * y, num_clusters_mn)
max_active_clusters = _get_max_active_clusters(cluster_shape_mn)
sm_count = _get_sm_count()
num_groups = len(problem_sizes)
compiled_func = cute.compile(
kernel_runner,
cute_ptr_of_tensor_of_problem_sizes,
cute_ptr_of_tensor_of_abc_ptrs,
cute_ptr_of_tensor_of_sfasfb_ptrs,
cute_ptr_of_tensor_of_tensormap,
cute_ptr_of_tensor_of_strides_abc,
sm_count,
total_num_clusters,
num_groups,
max_active_clusters
)
# Store compiled kernel in cache with problem_sizes as key
_compiled_kernel_cache[cache_key] = compiled_func
return compiled_func
def custom_kernel(data: input_t) -> output_t:
"""
Execute the block-scaled group GEMM kernel.
This is the main entry point called by the evaluation framework.
It converts PyTorch tensors to CuTe tensors, launches the kernel,
and returns the result.
Args:
data: Tuple of (abc_tensors, sfasfb_tensors, problem_sizes) where:
abc_tensors: list of tuples (a, b, c) where
a is torch.Tensor[float4e2m1fn_x2] of shape [m, k // 2, l]
b is torch.Tensor[float4e2m1fn_x2] of shape [n, k // 2, l]
c is torch.Tensor[float16] of shape [m, n, l]
sfasfb_tensors: list of tuples (sfa, sfb) where
sfa is torch.Tensor[float8_e4m3fnuz] of shape [m, k // 16, l]
sfb is torch.Tensor[float8_e4m3fnuz] of shape [n, k // 16, l]
problem_sizes: list of tuples (m, n, k, l)
each group has its own a, b, c, sfa, sfb with different m, n, k, l problem sizes
l should always be 1 for each group.
list size is the number of groups.
Returns:
list of c tensors where c is torch.Tensor[float16] of shape [m, n, l] for each group
"""
abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data
abc_ptrs = tuple((a.data_ptr(), b.data_ptr(), c.data_ptr()) for a, b, c in abc_tensors)
sfasfb_ptrs = tuple(
(sfa_reordered.data_ptr(), sfb_reordered.data_ptr())
for sfa_reordered, sfb_reordered in sfasfb_reordered_tensors
)
launch_cache_key = (
torch.cuda.current_device(),
make_problem_sizes_cache_key(problem_sizes),
abc_ptrs,
sfasfb_ptrs,
)
cached_launch = _launch_context_cache.get(launch_cache_key)
if cached_launch is None:
compiled_func = compile_kernel(problem_sizes)
tensor_of_problem_sizes, tensor_of_strides_abc = _get_static_metadata_tensors(
problem_sizes
)
# Cache pointer tensors to avoid per-call CPU work and H2D metadata upload.
# tensor_of_abc_ptrs: Shape (num_groups, 3) containing (a_ptr, b_ptr, c_ptr) per group
# tensor_of_sfasfb_ptrs: Shape (num_groups, 2) containing (sfa_ptr, sfb_ptr) per group
tensor_of_abc_ptrs, tensor_of_sfasfb_ptrs = _get_pointer_metadata_tensors(
abc_ptrs,
sfasfb_ptrs,
)
# Allocate device memory for tensormap descriptors
# Each cluster needs its own set of tensormaps (one for A, B, SFA, SFB, C)
# Shape: (sm_count, num_tensormaps=5, bytes_per_tensormap/8=16)
# sm_count because we are running persisent scheduler.
# Tensormaps are hardware descriptors used by TMA for efficient memory transfers
sm_count = _get_sm_count()
tensor_of_tensormap = _get_tensormap_workspace(sm_count)
# Create CuTe pointers to the metadata tensors that will be passed to the kernel
# These allow the GPU kernel to read problem sizes and tensor pointers
cute_ptr_of_tensor_of_abc_ptrs = make_ptr(
cutlass.Int64,
tensor_of_abc_ptrs.data_ptr(),
cute.AddressSpace.gmem,
assumed_align=16,
)
cute_ptr_of_tensor_of_sfasfb_ptrs = make_ptr(
cutlass.Int64,
tensor_of_sfasfb_ptrs.data_ptr(),
cute.AddressSpace.gmem,
assumed_align=16,
)
cute_ptr_of_tensor_of_problem_sizes = make_ptr(
cutlass.Int32,
tensor_of_problem_sizes.data_ptr(),
cute.AddressSpace.gmem,
assumed_align=16,
)
cute_ptr_of_tensor_of_tensormap = make_ptr(
cutlass.Int64,
tensor_of_tensormap.data_ptr(),
cute.AddressSpace.gmem,
assumed_align=16,
)
cute_ptr_of_tensor_of_strides_abc = make_ptr(
cutlass.Int32,
tensor_of_strides_abc.data_ptr(),
cute.AddressSpace.gmem,
assumed_align=16,
)
cached_launch = (
compiled_func,
cute_ptr_of_tensor_of_problem_sizes,
cute_ptr_of_tensor_of_abc_ptrs,
cute_ptr_of_tensor_of_sfasfb_ptrs,
cute_ptr_of_tensor_of_tensormap,
cute_ptr_of_tensor_of_strides_abc,
)
_launch_context_cache[launch_cache_key] = cached_launch
else:
(
compiled_func,
cute_ptr_of_tensor_of_problem_sizes,
cute_ptr_of_tensor_of_abc_ptrs,
cute_ptr_of_tensor_of_sfasfb_ptrs,
cute_ptr_of_tensor_of_tensormap,
cute_ptr_of_tensor_of_strides_abc,
) = cached_launch
# Launch the JIT-compiled GPU kernel with all prepared data
# The kernel will perform block-scaled group GEMM: C = A * SFA * B * SFB for all groups
compiled_func(
cute_ptr_of_tensor_of_problem_sizes, # Pointer to problem sizes array
cute_ptr_of_tensor_of_abc_ptrs, # Pointer to ABC tensor pointers array
cute_ptr_of_tensor_of_sfasfb_ptrs, # Pointer to scale factor pointers array
cute_ptr_of_tensor_of_tensormap, # Pointer to tensormap buffer
cute_ptr_of_tensor_of_strides_abc, # Pointer to tensor
)
num_groups = len(problem_sizes)
res = []
for i in range(num_groups):
res.append(abc_tensors[i][2])
return res
scrolls · 2475 lines total
Source code from GPU Mode and the KernelBot dataset · June 9 Researcher Reciprocity License v1.0
Changes from previous submission
Against this author's previous submission submission 477885.
⋯ 60 unchanged lines- and you’re targeting performance on SM100."""- cluster_shape_mn = (1, 1)+ cluster_shape_mn = (1, 2)class Sm100GroupedBlockScaledGemmKernel:# Size of smem we reserved for mbarrier, tensor memory management and tensormap update⋯ 581 unchanged linestensormap_manager.fence_tensormap_update(tensormap_b_gmem_ptr)tensormap_manager.fence_tensormap_update(tensormap_sfa_gmem_ptr)tensormap_manager.fence_tensormap_update(tensormap_sfb_gmem_ptr)+ tma_desc_ptr_a = tensormap_manager.get_tensormap_ptr(+ tensormap_a_gmem_ptr,+ cute.AddressSpace.generic,+ )+ tma_desc_ptr_b = tensormap_manager.get_tensormap_ptr(+ tensormap_b_gmem_ptr,+ cute.AddressSpace.generic,+ )+ tma_desc_ptr_sfa = tensormap_manager.get_tensormap_ptr(+ tensormap_sfa_gmem_ptr,+ cute.AddressSpace.generic,+ )+ tma_desc_ptr_sfb = tensormap_manager.get_tensormap_ptr(+ tensormap_sfb_gmem_ptr,+ cute.AddressSpace.generic,+ )## Tma load loop#⋯ 10 unchanged linestAsA[(None, ab_producer_state.index)],tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),mcast_mask=a_full_mcast_mask,- tma_desc_ptr=tensormap_manager.get_tensormap_ptr(- tensormap_a_gmem_ptr,- cute.AddressSpace.generic,- ),+ tma_desc_ptr=tma_desc_ptr_a,)cute.copy(tma_atom_b,⋯ 1 unchanged linestBsB[(None, ab_producer_state.index)],tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),mcast_mask=b_full_mcast_mask,- tma_desc_ptr=tensormap_manager.get_tensormap_ptr(- tensormap_b_gmem_ptr,- cute.AddressSpace.generic,- ),+ tma_desc_ptr=tma_desc_ptr_b,)cute.copy(tma_atom_sfa,⋯ 1 unchanged linestAsSFA[(None, ab_producer_state.index)],tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),mcast_mask=sfa_full_mcast_mask,- tma_desc_ptr=tensormap_manager.get_tensormap_ptr(- tensormap_sfa_gmem_ptr,- cute.AddressSpace.generic,- ),+ tma_desc_ptr=tma_desc_ptr_sfa,)cute.copy(tma_atom_sfb,⋯ 1 unchanged linestBsSFB[(None, ab_producer_state.index)],tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),mcast_mask=sfb_full_mcast_mask,- tma_desc_ptr=tensormap_manager.get_tensormap_ptr(- tensormap_sfb_gmem_ptr,- cute.AddressSpace.generic,- ),+ tma_desc_ptr=tma_desc_ptr_sfb,)# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1⋯ 210 unchanged linestCrB[kblock_coord],tCtAcc,)+ # Enable accumulate after the first kblock only.+ if kblock_idx == 0:+ tiled_mma.set(tcgen05.Field.ACCUMULATE, True)- # Enable accumulate on tCtAcc after first kblock- tiled_mma.set(tcgen05.Field.ACCUMULATE, True)-# Async arrive AB buffer emptyab_pipeline.consumer_release(ab_consumer_state)⋯ 189 unchanged lines#subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt+ tma_desc_ptr_c = tensormap_manager.get_tensormap_ptr(+ tensormap_c_gmem_ptr,+ cute.AddressSpace.generic,+ )+ tRS_rAcc = tiled_copy_r2s.retile(tTR_rAcc)for subtile_idx in range(subtile_cnt):## Load accumulator from tensor memory buffer to register⋯ 4 unchanged lines## Convert to C type#- acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load()+ acc_vec = tRS_rAcc.load()tRS_rC.store(acc_vec.to(c_dtype))#⋯ 20 unchanged linestma_atom_c,bSG_sC[(None, c_buffer)],bSG_gC[(None, subtile_idx)],- tma_desc_ptr=tensormap_manager.get_tensormap_ptr(- tensormap_c_gmem_ptr,- cute.AddressSpace.generic,- ),+ tma_desc_ptr=tma_desc_ptr_c,)# Fence and barrier to make sure shared memory store is visible to TMA storec_pipeline.producer_commit()⋯ 419 unchanged lines# Compute grid size (For first problem)# total_num_clusters = 352- # self.cluster_shape_mn = (1, 1)+ # self.cluster_shape_mn = (1, 2)# max_active_clusters = 148self.tile_sched_params, grid = self._compute_grid(total_num_clusters, self.cluster_shape_mn, max_active_clusters⋯ 558 unchanged lines_compiled_kernel_cache = {}_static_metadata_cache = {}_tensormap_workspace_cache = {}+ _pointer_metadata_cache = {}+ _launch_context_cache = {}_hardware_info = None_sm_count_cache = None_max_active_clusters_cache = {}⋯ 54 unchanged linesreturn tensor_of_problem_sizes, tensor_of_strides_abc+ def _get_pointer_metadata_tensors(+ abc_ptrs,+ sfasfb_ptrs,+ ) -> Tuple[torch.Tensor, torch.Tensor]:+ device_idx = torch.cuda.current_device()+ cache_key = (+ device_idx,+ tuple(abc_ptrs),+ tuple(sfasfb_ptrs),+ )+ cached = _pointer_metadata_cache.get(cache_key)+ if cached is not None:+ return cached++ tensor_of_abc_ptrs = torch.tensor(abc_ptrs, dtype=torch.int64, device="cuda")+ tensor_of_sfasfb_ptrs = torch.tensor(sfasfb_ptrs, dtype=torch.int64, device="cuda")+ _pointer_metadata_cache[cache_key] = (tensor_of_abc_ptrs, tensor_of_sfasfb_ptrs)+ return tensor_of_abc_ptrs, tensor_of_sfasfb_ptrs++def _get_tensormap_workspace(sm_count: int) -> torch.Tensor:device_idx = torch.cuda.current_device()cache_key = (device_idx, sm_count)⋯ 62 unchanged lines# Compute the tile shape for each CUDA Thread Block (CTA)# cta_tile_shape_mn: [M_tile, N_tile] = [128, 128] for this kernelcta_tile_shape_mn = [128, mma_tiler_mnk[1]]- # cluster_tile_shape_mn: Total tile shape per cluster (same as CTA since cluster is 1x1)+ # cluster_tile_shape_mn: Total tile shape per clustercluster_tile_shape_mn = tuple(- x * y for x, y in zip(cta_tile_shape_mn, (1, 1))+ x * y for x, y in zip(cta_tile_shape_mn, cluster_shape_mn))# Compute total number of cluster tiles needed across all groups⋯ 57 unchanged lines"""abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data- compiled_func = compile_kernel(problem_sizes)-- # Extract raw data pointers from all input tensors for each group- # These will be passed to the GPU kernel to access the actual tensor data- abc_ptrs = []- sfasfb_ptrs = []- for i, ((a, b, c), (sfa_reordered, sfb_reordered), (m, n, k, l)) in enumerate(zip(abc_tensors, sfasfb_reordered_tensors, problem_sizes)):- # Store pointers to A, B, and C matrices for this group- abc_ptrs.append((a.data_ptr(), b.data_ptr(), c.data_ptr()))- # Store pointers to scale factor tensors for this group- sfasfb_ptrs.append((sfa_reordered.data_ptr(), sfb_reordered.data_ptr()))- tensor_of_problem_sizes, tensor_of_strides_abc = _get_static_metadata_tensors(- problem_sizes+ abc_ptrs = tuple((a.data_ptr(), b.data_ptr(), c.data_ptr()) for a, b, c in abc_tensors)+ sfasfb_ptrs = tuple(+ (sfa_reordered.data_ptr(), sfb_reordered.data_ptr())+ for sfa_reordered, sfb_reordered in sfasfb_reordered_tensors)+ launch_cache_key = (+ torch.cuda.current_device(),+ make_problem_sizes_cache_key(problem_sizes),+ abc_ptrs,+ sfasfb_ptrs,+ )+ cached_launch = _launch_context_cache.get(launch_cache_key)+ if cached_launch is None:+ compiled_func = compile_kernel(problem_sizes)- # Create torch tensors to store data pointers for all groups- # These allow the GPU kernel to dynamically access different tensors per group- # tensor_of_abc_ptrs: Shape (num_groups, 3) containing (a_ptr, b_ptr, c_ptr) per group- # tensor_of_sfasfb_ptrs: Shape (num_groups, 2) containing (sfa_ptr, sfb_ptr) per group- tensor_of_abc_ptrs = torch.tensor(abc_ptrs, dtype=torch.int64, device="cuda")- tensor_of_sfasfb_ptrs = torch.tensor(sfasfb_ptrs, dtype=torch.int64, device="cuda")+ tensor_of_problem_sizes, tensor_of_strides_abc = _get_static_metadata_tensors(+ problem_sizes+ )- # Allocate device memory for tensormap descriptors- # Each cluster needs its own set of tensormaps (one for A, B, SFA, SFB, C)- # Shape: (sm_count, num_tensormaps=5, bytes_per_tensormap/8=16)- # sm_count because we are running persisent scheduler.- # Tensormaps are hardware descriptors used by TMA for efficient memory transfers- sm_count = _get_sm_count()- tensor_of_tensormap = _get_tensormap_workspace(sm_count)+ # Cache pointer tensors to avoid per-call CPU work and H2D metadata upload.+ # tensor_of_abc_ptrs: Shape (num_groups, 3) containing (a_ptr, b_ptr, c_ptr) per group+ # tensor_of_sfasfb_ptrs: Shape (num_groups, 2) containing (sfa_ptr, sfb_ptr) per group+ tensor_of_abc_ptrs, tensor_of_sfasfb_ptrs = _get_pointer_metadata_tensors(+ abc_ptrs,+ sfasfb_ptrs,+ )- # Create CuTe pointers to the metadata tensors that will be passed to the kernel- # These allow the GPU kernel to read problem sizes and tensor pointers- cute_ptr_of_tensor_of_abc_ptrs = make_ptr(- cutlass.Int64,- tensor_of_abc_ptrs.data_ptr(),- cute.AddressSpace.gmem,- assumed_align=16,- )- cute_ptr_of_tensor_of_sfasfb_ptrs = make_ptr(- cutlass.Int64,- tensor_of_sfasfb_ptrs.data_ptr(),- cute.AddressSpace.gmem,- assumed_align=16,- )- cute_ptr_of_tensor_of_problem_sizes = make_ptr(- cutlass.Int32,- tensor_of_problem_sizes.data_ptr(),- cute.AddressSpace.gmem,- assumed_align=16,- )- cute_ptr_of_tensor_of_tensormap = make_ptr(- cutlass.Int64,- tensor_of_tensormap.data_ptr(),- cute.AddressSpace.gmem,- assumed_align=16,- )- cute_ptr_of_tensor_of_strides_abc = make_ptr(- cutlass.Int32,- tensor_of_strides_abc.data_ptr(),- cute.AddressSpace.gmem,- assumed_align=16,- )+ # Allocate device memory for tensormap descriptors+ # Each cluster needs its own set of tensormaps (one for A, B, SFA, SFB, C)+ # Shape: (sm_count, num_tensormaps=5, bytes_per_tensormap/8=16)+ # sm_count because we are running persisent scheduler.+ # Tensormaps are hardware descriptors used by TMA for efficient memory transfers+ sm_count = _get_sm_count()+ tensor_of_tensormap = _get_tensormap_workspace(sm_count)+ # Create CuTe pointers to the metadata tensors that will be passed to the kernel+ # These allow the GPU kernel to read problem sizes and tensor pointers+ cute_ptr_of_tensor_of_abc_ptrs = make_ptr(+ cutlass.Int64,+ tensor_of_abc_ptrs.data_ptr(),+ cute.AddressSpace.gmem,+ assumed_align=16,+ )+ cute_ptr_of_tensor_of_sfasfb_ptrs = make_ptr(+ cutlass.Int64,+ tensor_of_sfasfb_ptrs.data_ptr(),+ cute.AddressSpace.gmem,+ assumed_align=16,+ )+ cute_ptr_of_tensor_of_problem_sizes = make_ptr(+ cutlass.Int32,+ tensor_of_problem_sizes.data_ptr(),+ cute.AddressSpace.gmem,+ assumed_align=16,+ )+ cute_ptr_of_tensor_of_tensormap = make_ptr(+ cutlass.Int64,+ tensor_of_tensormap.data_ptr(),+ cute.AddressSpace.gmem,+ assumed_align=16,+ )+ cute_ptr_of_tensor_of_strides_abc = make_ptr(+ cutlass.Int32,+ tensor_of_strides_abc.data_ptr(),+ cute.AddressSpace.gmem,+ assumed_align=16,+ )++ cached_launch = (+ compiled_func,+ cute_ptr_of_tensor_of_problem_sizes,+ cute_ptr_of_tensor_of_abc_ptrs,+ cute_ptr_of_tensor_of_sfasfb_ptrs,+ cute_ptr_of_tensor_of_tensormap,+ cute_ptr_of_tensor_of_strides_abc,+ )+ _launch_context_cache[launch_cache_key] = cached_launch+ else:+ (+ compiled_func,+ cute_ptr_of_tensor_of_problem_sizes,+ cute_ptr_of_tensor_of_abc_ptrs,+ cute_ptr_of_tensor_of_sfasfb_ptrs,+ cute_ptr_of_tensor_of_tensormap,+ cute_ptr_of_tensor_of_strides_abc,+ ) = cached_launch+# Launch the JIT-compiled GPU kernel with all prepared data# The kernel will perform block-scaled group GEMM: C = A * SFA * B * SFB for all groupscompiled_func(
scrolls · 336 diff lines total
Best evidence level for this revision: reported
JSON