submission 506091
Inoday · python · License unknown
Use it
Vendorable · source mirrored · license unknownView source →
No package. Vendor the mirrored source: 2787 lines, June 9 Researcher Reciprocity License v1.0.
submission_1.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-nvfp4-group-gemm-506091?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:ed7ce61be5352edb9e8853acd5c51c6012ff05a78f43cfe1fb9e534256592689
license declaredunknown
license concludedunknown
authorsInoday
imported2026-08-15
Techniques
Extracted from the mirrored source by pattern, never inferred. Each row cites its line.
fused-epilogue
- Computing epilogue subtilembarrier
self.epilog_sync_barrier = pipeline.NamedBarrier(persistent-kernel
warp-specialized persistent kernel.shared-memory
self.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_1.py2787 lines
#!POPCORN leaderboard nvfp4_group_gemm
#!POPCORN gpu NVIDIA
# Thanks to Simon Veitner for sharing his TVM FFI compilation setup https://gist.github.com/simveit/baec5535ff6c5bd594c6f897d38075c8
import subprocess
import sys
for pkg in ["apache-tvm-ffi", "torch-c-dlpack-ext"]:
subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", pkg])
import argparse
from math import prod
import functools
from inspect import isclass
from typing import List, Tuple, Type, Union
import cuda.bindings.driver as cuda
import torch
import cutlass
import cutlass.cute as cute
import cutlass.pipeline as pipeline
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.cute.nvgpu import cpasync, tcgen05
from cutlass.cutlass_dsl import const_expr
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
from task import input_t, output_t
# 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
# TMA L2 cache-hint policies used by prior top submissions.
_TMA_CACHE_EVICT_NORMAL = 0x1000000000000000
_TMA_CACHE_EVICT_FIRST = 0x12F0000000000000
_TMA_CACHE_EVICT_LAST = 0x14F0000000000000
# Supported benchmark families where only M is dynamic.
# family_id -> (group_count, n_const, k_const, mma_tiler_mn, enable_cache_hints)
_BENCHMARK_FAMILY_SPECS: dict[int, tuple[int, int, int, tuple[int, int], bool]] = {
0: (8, 4096, 7168, (128, 128), True),
1: (8, 7168, 2048, (128, 256), False),
2: (2, 3072, 4096, (128, 128), False),
3: (2, 4096, 1536, (128, 128), False),
}
_BENCHMARK_FAMILY_BY_KEY: dict[tuple[int, int, int], int] = {
(spec[0], spec[1], spec[2]): family_id
for family_id, spec in _BENCHMARK_FAMILY_SPECS.items()
}
class Sm100GroupedBlockScaledGemmKernel:
"""This example demonstrates an implementation of grouped blockscaled GEMM using a TMA plus Blackwell SM100 TensorCore
warp-specialized persistent kernel.
:param sf_vec_size: Scalefactor vector size.
:type sf_vec_size: int
:param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N)
:type mma_tiler_mn: Tuple[int, int]
:param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing
:type cluster_shape_mn: Tuple[int, int]
:note: In current version, A and B tensors must have the same data type
- i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported
:note: Supported combinations of A/B data types, SF data typs and SF vector size:
- MXF8: A/B: Float8E5M2/Float8E4M3FN + SF: Float8E8M0FNU + sf_vec_size: 32
- MXF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU + sf_vec_size: 32
- NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN + sf_vec_size: 16
:note: Supported accumulator data types:
- Float32
:note: Supported C data types:
- Float32
- Float16/BFloat16
- Float8E4M3FN/Float8E5M2
:note: Constraints:
- MMA tiler M must be 128 or 256 (use_2cta_instrs)
- MMA tiler N must be 128/256
- Cluster shape M must be multiple of 2 if Mma tiler M is 256
- Cluster shape M/N must be positive and power of 2, total cluster size <= 16
- Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors
"""
def __init__(
self,
sf_vec_size: int,
mma_tiler_mn: Tuple[int, int],
cluster_shape_mn: Tuple[int, int],
enable_cache_hints: bool,
ab_stage_cap: int | None = None,
problem_shape_n_const: int | None = None,
problem_shape_k_const: int | None = None,
):
"""Initializes the configuration for a Blackwell grouped blockscaled GEMM kernel.
Besides configurations for dense persistent blockscaled GEMM, there is an extra config specific to grouped blockscaled GEMM:
:param sf_vec_size: Scalefactor vector size.
:type sf_vec_size: int
:param mma_tiler_mn: tuple (M, N) shape of the MMA instruction.
:type mma_tiler_mn: tuple[int, int]
:param cluster_shape_mn: tuple (ClusterM, ClusterN) shape of the cluster.
:type cluster_shape_mn: tuple[int, int]
"""
self.acc_dtype = cutlass.Float32
self.sf_vec_size = sf_vec_size
self.use_2cta_instrs = mma_tiler_mn[0] == 256
self.cluster_shape_mn = cluster_shape_mn
self.enable_cache_hints = enable_cache_hints
self.ab_stage_cap = ab_stage_cap
self.problem_shape_n_const = problem_shape_n_const
self.problem_shape_k_const = problem_shape_k_const
# 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.occupancy = 1
# Set specialized warp ids
self.epilog_warp_id = (
0,
1,
2,
3,
)
self.mma_warp_id = 4
self.tma_warp_id = 5
self.threads_per_cta = 32 * len(
(self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
)
# 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,
)
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
SM100_TMEM_CAPACITY_COLUMNS = 512
self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS
# Set up configurations that dependent on gemm inputs.
def _setup_attributes(self):
"""Set up configurations that are dependent on GEMM inputs
This method configures various attributes based on the input tensor properties
(data types, leading dimensions) and kernel settings:
- Configuring tiled MMA
- Computing MMA/cluster/tile shapes
- Computing cluster layout
- Computing multicast CTAs for A/B/SFA/SFB
- Computing epilogue subtile
- Setting up A/B/SFA/SFB/C stage counts in shared memory
- Computing A/B/SFA/SFB/C shared memory layout
- Checking reserved smem bytes size capacity for mbar, tensor memory management and tensormap updates utilization
"""
# 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),
)
tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.a_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
self.cta_group,
self.mma_inst_shape_mn,
)
tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.a_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
cute.nvgpu.tcgen05.CtaGroup.ONE,
self.mma_inst_shape_mn_sfb,
)
# Compute mma/cluster/tile shapes
mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
mma_inst_tile_k = 4
self.mma_tiler = (
self.mma_inst_shape_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))
)
# 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
# Compute epilogue subtile
self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
self.cta_tile_shape_mnk,
self.use_2cta_instrs,
self.c_layout,
self.c_dtype,
)
# Setup A/B stage count in shared memory and ACC stage count in tensor memory
self.num_acc_stage, self.num_ab_stage = self._compute_stages(
tiled_mma,
self.mma_tiler,
self.a_dtype,
self.b_dtype,
self.epi_tile,
self.c_dtype,
self.c_layout,
self.sf_dtype,
self.sf_vec_size,
self.smem_capacity,
self.occupancy,
)
# Family-specific override for AB stages from benchmark-family selector.
if self.ab_stage_cap is not None:
self.num_ab_stage = max(1, self.ab_stage_cap)
# Overlap and ping-pong accumulator buffers when ACC pipeline has a single stage.
self.overlapping_accum = self.num_acc_stage == 1
self.epi_tile_n = cute.size(self.epi_tile[1])
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.mma_tiler_sfb[1] // sf_atom_mn
) * mma_inst_tile_k
self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols
num_subtiles_in_overlap_region = (
self.num_sf_tmem_cols + self.epi_tile_n - 1
) // self.epi_tile_n
self.iter_acc_early_release_in_epilogue = (
num_subtiles_in_overlap_region - 1
)
# Compute A/B/SFA/SFB/C shared memory layout
self.a_smem_layout_staged = sm100_utils.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(
tiled_mma,
self.mma_tiler,
self.b_dtype,
self.num_ab_stage,
)
self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
tiled_mma,
self.mma_tiler,
self.sf_vec_size,
self.num_ab_stage,
)
self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
tiled_mma,
self.mma_tiler,
self.sf_vec_size,
self.num_ab_stage,
)
mbar_smem_bytes = self._get_mbar_smem_bytes(
num_acc_stage=self.num_acc_stage,
num_ab_stage=self.num_ab_stage,
)
# Use utils.TensorMapUpdateMode.SMEM by default
tensormap_smem_bytes = (
Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap
* Sm100GroupedBlockScaledGemmKernel.num_tensormaps
)
if (
mbar_smem_bytes
+ tensormap_smem_bytes
+ Sm100GroupedBlockScaledGemmKernel.tensor_memory_management_bytes
> self.reserved_smem_bytes
):
raise ValueError(
f"smem consumption for mbar and tensormap {mbar_smem_bytes + tensormap_smem_bytes} exceeds the "
f"reserved smem bytes {self.reserved_smem_bytes}"
)
@cute.jit
def __call__(
self,
group_count: cutlass.Constexpr[int],
problem_shape_mnkl: cute.Tensor,
group_runtime_meta: cute.Tensor,
use_2cta_instrs: cutlass.Constexpr[bool],
strides_abc: cute.Tensor,
tensor_address_abc: cute.Tensor,
tensor_address_sfasfb: cute.Tensor,
total_num_clusters: cutlass.Constexpr[int],
tensormap_cute_tensor: cute.Tensor,
max_active_clusters: cutlass.Constexpr[int],
):
"""Execute the GEMM operation in steps:
- Setup static attributes before smem/grid/tma computation
- Setup TMA load/store atoms and tensors
- Compute grid size with regard to hardware constraints
- Define shared storage for kernel
- Launch the kernel synchronously
For grouped GEMM, tensor shapes, tensor strides, and tensor address are all provided
by different tensors in global memory. Initial tensors for TMA setup are created
internally with fake shapes.
:param group_count: The number of GEMM groups.
:type group_count: cutlass.Constexpr[int]
:param problem_shape_mnkl: Tensor (group_count, 4):(4, 1) containing (M, N, K, L) for each group.
:type problem_shape_mnkl: cute.Tensor
:param group_runtime_meta: Tensor (group_count, 4):(4, 1) containing runtime metadata per group:
[M, cluster_tile_count_mn, _, _].
:type group_runtime_meta: cute.Tensor
:param use_2cta_instrs: Whether to use 2-CTA instructions.
:type use_2cta_instrs: cutlass.Constexpr[bool]
:param strides_abc: Tensor (group_count, 3, 2):(6, 2, 1) containing strides for A, B, C.
:type strides_abc: cute.Tensor
:param tensor_address_abc: Tensor (group_count, 3):(3, 1) containing addresses for A, B, C.
:type tensor_address_abc: cute.Tensor
:param tensor_address_sfasfb: Tensor (group_count, 2):(2, 1) containing addresses for SFA, SFB.
:type tensor_address_sfasfb: cute.Tensor
:param total_num_clusters: Total number of clusters needed for all groups.
:type total_num_clusters: cutlass.Constexpr[int]
:param tensormap_cute_tensor: Tensor (num_tensormap_buffers, 5, 16):(80, 16, 1) for tensormaps.
:type tensormap_cute_tensor: cute.Tensor
:param max_active_clusters: Maximum number of active clusters.
:type max_active_clusters: cutlass.Constexpr[int]
"""
# Use fake shape for initial TMA descriptor and atom setup
# The real TMA desc and atom will be updated during kernel execution.
min_shape = (cutlass.Int32(64), cutlass.Int32(64), cutlass.Int32(1))
# Create initial_a tensor with fake shape and null pointer
# Layout: (M, K, L) with K-major strides (K, 1, M*K)
initial_a = cute.make_tensor(
cute.make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
cute.make_layout(
(min_shape[0], min_shape[1], min_shape[2]),
stride=(
min_shape[1],
1,
min_shape[0] * min_shape[1],
),
),
)
# Create initial_b tensor with fake shape and null pointer
# Layout: (N, K, L) with K-major strides (K, 1, N*K)
initial_b = cute.make_tensor(
cute.make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
cute.make_layout(
(min_shape[0], min_shape[1], min_shape[2]),
stride=(
min_shape[1],
1,
min_shape[0] * min_shape[1],
),
),
)
# Create initial_c tensor with fake shape and null pointer
# Layout: (M, N, L) with N-major (row-major) strides (N, 1, M*N)
# 32B alignment for STG.256 (vs 16B for TMA/STG.128)
initial_c = cute.make_tensor(
cute.make_ptr(c_dtype, 0, cute.AddressSpace.gmem, assumed_align=32), # Changed from 16 to 32
cute.make_layout(
(min_shape[0], min_shape[1], min_shape[2]),
stride=(
min_shape[1],
1,
min_shape[0] * min_shape[1],
),
),
)
# Set data types from initial tensors and module-level constants
self.a_dtype = initial_a.element_type
self.b_dtype = initial_b.element_type
self.sf_dtype = sf_dtype
self.c_dtype = initial_c.element_type
self.a_major_mode = utils.LayoutEnum.from_tensor(initial_a).mma_major_mode()
self.b_major_mode = utils.LayoutEnum.from_tensor(initial_b).mma_major_mode()
self.c_layout = utils.LayoutEnum.from_tensor(initial_c)
# Setup attributes that dependent on gemm inputs
self._setup_attributes()
# 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, self.sf_vec_size
)
initial_sfa = cute.make_tensor(
cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
sfa_layout,
)
# ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
initial_b.shape, self.sf_vec_size
)
initial_sfb = cute.make_tensor(
cute.make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=16),
sfb_layout,
)
tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.a_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
self.cta_group,
self.mma_inst_shape_mn,
)
tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.a_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
cute.nvgpu.tcgen05.CtaGroup.ONE,
self.mma_inst_shape_mn_sfb,
)
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(
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(self.a_dtype, a_smem_layout)
b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout)
sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)
sfb_copy_size = cute.size_in_bytes(self.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,
# self.epi_tile,
# )
# Compute grid size
self.tile_sched_params, grid = self._compute_grid(
total_num_clusters, self.cluster_shape_mn, max_active_clusters
)
self.buffer_align_bytes = 1024
self.size_tensormap_in_i64 = (
Sm100GroupedBlockScaledGemmKernel.num_tensormaps
* Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap
// 8
)
# Define shared storage for kernel
@cute.struct
class SharedStorage:
tensormap_buffer: cute.struct.MemRange[
cutlass.Int64, self.size_tensormap_in_i64
]
# Cache grouped problem metadata in SMEM so group scheduler search avoids repeated GMEM reads.
problem_shape_mnkl_cache: cute.struct.MemRange[
cutlass.Int32, group_count * 4
]
# Fast-path caches for dynamic-M/static-NK scheduling.
problem_shape_m_cache: cute.struct.MemRange[cutlass.Int32, group_count]
cluster_tile_count_mn_cache: cute.struct.MemRange[cutlass.Int32, group_count]
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: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (MMA, MMA_M, MMA_K, STAGE)
sA: cute.struct.Align[
cute.struct.MemRange[
self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)
],
self.buffer_align_bytes,
]
# (MMA, MMA_N, MMA_K, STAGE)
sB: cute.struct.Align[
cute.struct.MemRange[
self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)
],
self.buffer_align_bytes,
]
# (MMA, MMA_M, MMA_K, STAGE)
sSFA: cute.struct.Align[
cute.struct.MemRange[
self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)
],
self.buffer_align_bytes,
]
# (MMA, MMA_N, MMA_K, STAGE)
sSFB: cute.struct.Align[
cute.struct.MemRange[
self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)
],
self.buffer_align_bytes,
]
self.shared_storage = SharedStorage
# Launch the kernel synchronously
self.kernel(
tiled_mma,
tiled_mma_sfb,
tma_atom_a,
tma_tensor_a,
tma_atom_b,
tma_tensor_b,
tma_atom_sfa,
tma_tensor_sfa,
tma_atom_sfb,
tma_tensor_sfb,
initial_c,
self.cluster_layout_vmnk,
self.cluster_layout_sfb_vmnk,
self.a_smem_layout_staged,
self.b_smem_layout_staged,
self.sfa_smem_layout_staged,
self.sfb_smem_layout_staged,
self.epi_tile,
self.tile_sched_params,
group_count,
problem_shape_mnkl,
group_runtime_meta,
use_2cta_instrs,
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
# GPU device kernel
@cute.kernel
def 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,
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,
epi_tile: cute.Tile,
tile_sched_params: utils.PersistentTileSchedulerParams,
group_count: cutlass.Constexpr,
problem_sizes_mnkl: cute.Tensor,
group_runtime_meta: cute.Tensor,
use_2cta_instrs: cutlass.Constexpr[bool],
strides_abc: cute.Tensor,
ptrs_abc: cute.Tensor,
ptrs_sfasfb: cute.Tensor,
tensormaps: cute.Tensor,
):
"""
GPU device kernel performing the grouped GEMM computation.
"""
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
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)
#
# Setup cta/thread coordinates
#
# Coords inside cluster
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
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()
#
# Alloc and init: tensormap buffer, a+b full/empty, accumulator full/empty, tensor memory dealloc barrier
#
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
)
problem_shape_mnkl_smem = storage.problem_shape_mnkl_cache.get_tensor(
cute.make_layout((group_count, 4), stride=(4, 1))
)
problem_shape_m_smem = storage.problem_shape_m_cache.get_tensor(
cute.make_layout((group_count,), stride=(1,))
)
cluster_tile_count_mn_smem = storage.cluster_tile_count_mn_cache.get_tensor(
cute.make_layout((group_count,), stride=(1,))
)
problem_shape_mnkl_smem_flat = cute.flatten(problem_shape_mnkl_smem)
problem_shape_mnkl_gmem_flat = cute.flatten(problem_sizes_mnkl)
group_runtime_meta_flat = cute.flatten(group_runtime_meta)
problem_shape_n_const = cutlass.Int32(0)
problem_shape_k_const = cutlass.Int32(0)
cluster_tile_count_n_const = cutlass.Int32(0)
cluster_tile_count_k_const = cutlass.Int32(0)
if cutlass.const_expr(
not (
self.problem_shape_n_const is None
or self.problem_shape_k_const is None
)
):
problem_shape_n_const = cutlass.Int32(self.problem_shape_n_const)
problem_shape_k_const = cutlass.Int32(self.problem_shape_k_const)
cluster_tile_count_n_const = (
problem_shape_n_const + self.cluster_tile_shape_mnk[1] - 1
) // self.cluster_tile_shape_mnk[1]
cluster_tile_count_k_const = (
problem_shape_k_const + self.cluster_tile_shape_mnk[2] - 1
) // self.cluster_tile_shape_mnk[2]
# Copy runtime metadata densely as int32 for coalesced global loads.
for i in cutlass.range(tidx, group_count * 4, self.threads_per_cta):
problem_shape_mnkl_smem_flat[i] = group_runtime_meta_flat[i]
cute.arch.sync_threads()
# Extract only M and precomputed cluster_tile_count_mn from SMEM cache.
for i in cutlass.range(tidx, group_count, self.threads_per_cta):
problem_shape_m_smem[(i,)] = problem_shape_mnkl_smem[(i, 0)]
cluster_tile_count_mn_smem[(i,)] = problem_shape_mnkl_smem[(i, 1)]
cute.arch.sync_threads()
else:
for i in cutlass.range(tidx, group_count * 4, self.threads_per_cta):
problem_shape_mnkl_smem_flat[i] = problem_shape_mnkl_gmem_flat[i]
cute.arch.sync_threads()
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 cutlass.const_expr(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 cutlass.const_expr(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
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
#
# Setup smem tensor A/B/SFA/SFB/C
#
# (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
#
# (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
#
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)
#
# 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])
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)
)
# Pack two logical accumulator stages into TMEM with SFA/SFB overlap.
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
#
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
#
# Get tensormap buffer address
#
grid_dim = cute.arch.grid_dim()
tensormap_workspace_idx = (
bidz * grid_dim[1] * grid_dim[0] + bidy * grid_dim[0] + bidx
)
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
)
#
# Persistent tile scheduling loop
#
tile_sched = utils.StaticPersistentTileScheduler.create(
tile_sched_params, cute.arch.block_idx(), grid_dim
)
group_gemm_ts_helper = None
if cutlass.const_expr(
self.problem_shape_n_const is None or self.problem_shape_k_const is None
):
group_gemm_ts_helper = utils.GroupedGemmTileSchedulerHelper(
group_count,
tile_sched_params,
self.cluster_tile_shape_mnk,
utils.create_initial_search_state(),
)
initial_work_tile_info = tile_sched.initial_work_tile_info()
#
# Specialized TMA load warp
#
if warp_idx == self.tma_warp_id and initial_work_tile_info.is_valid_tile:
tensormap_init_done = cutlass.Boolean(False)
# group index of last tile
last_group_idx = cutlass.Int32(-1)
work_tile = initial_work_tile_info
ab_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_ab_stage
)
while work_tile.is_valid_tile:
found_group = cutlass.Boolean(True)
cur_group_idx = cutlass.Int32(0)
cta_tile_idx_m = cutlass.Int32(0)
cta_tile_idx_n = cutlass.Int32(0)
cur_k_tile_cnt = cutlass.Int32(0)
problem_shape_m_group = cutlass.Int32(0)
problem_shape_n_group = cutlass.Int32(0)
problem_shape_k_group = cutlass.Int32(0)
if cutlass.const_expr(
not (
self.problem_shape_n_const is None
or self.problem_shape_k_const is None
)
):
(
found_group,
cur_group_idx,
cta_tile_idx_m,
cta_tile_idx_n,
problem_shape_m_group,
) = self.lookup_group_tile_info_dynamic_m(
work_tile.tile_idx,
group_count,
problem_shape_m_smem,
cluster_tile_count_mn_smem,
cluster_tile_count_n_const,
)
cur_k_tile_cnt = cluster_tile_count_k_const
problem_shape_n_group = problem_shape_n_const
problem_shape_k_group = problem_shape_k_const
else:
grouped_gemm_cta_tile_info = group_gemm_ts_helper.delinearize_z(
work_tile.tile_idx, problem_shape_mnkl_smem
)
cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k
cur_group_idx = grouped_gemm_cta_tile_info.group_idx
cta_tile_idx_m = grouped_gemm_cta_tile_info.cta_tile_idx_m
cta_tile_idx_n = grouped_gemm_cta_tile_info.cta_tile_idx_n
problem_shape_m_group = grouped_gemm_cta_tile_info.problem_shape_m
problem_shape_n_group = grouped_gemm_cta_tile_info.problem_shape_n
problem_shape_k_group = grouped_gemm_cta_tile_info.problem_shape_k
is_group_changed = cur_group_idx != last_group_idx
if found_group:
# 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,
self.a_dtype,
(
problem_shape_m_group,
problem_shape_n_group,
problem_shape_k_group,
),
strides_abc,
ptrs_abc,
0, # 0 for tensor A
)
real_tensor_b = self.make_tensor_abc_for_tensormap_update(
cur_group_idx,
self.b_dtype,
(
problem_shape_m_group,
problem_shape_n_group,
problem_shape_k_group,
),
strides_abc,
ptrs_abc,
1, # 1 for tensor B
)
real_tensor_sfa = self.make_tensor_sfasfb_for_tensormap_update(
cur_group_idx,
self.sf_dtype,
(
problem_shape_m_group,
problem_shape_n_group,
problem_shape_k_group,
),
ptrs_sfasfb,
0, # 0 for tensor SFA
)
real_tensor_sfb = self.make_tensor_sfasfb_for_tensormap_update(
cur_group_idx,
self.sf_dtype,
(
problem_shape_m_group,
problem_shape_n_group,
problem_shape_k_group,
),
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,
),
)
mma_tile_coord_mnl = (
cta_tile_idx_m // cute.size(tiled_mma.thr_id.shape),
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
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)
# Pick asymmetric TMA cache policies based on tile-shape dominance.
# If N-side has more cluster tiles, A/SFA are reused more across N-tiles,
# so keep A/SFA longer (EVICT_LAST) and B/SFB (EVICT_FIRST).
m_cluster_tiles = (
problem_shape_m_group
+ self.cluster_tile_shape_mnk[0]
- 1
) // self.cluster_tile_shape_mnk[0]
if cutlass.const_expr(
not (
self.problem_shape_n_const is None
or self.problem_shape_k_const is None
)
):
n_cluster_tiles = cluster_tile_count_n_const
else:
n_cluster_tiles = (
problem_shape_n_group
+ self.cluster_tile_shape_mnk[1]
- 1
) // self.cluster_tile_shape_mnk[1]
# DSL requirement: initialize values before dynamic control flow.
tma_cache_policy_a_sfa = cutlass.Int64(
cutlass.Int64(_TMA_CACHE_EVICT_NORMAL).ir_value()
)
tma_cache_policy_b_sfb = cutlass.Int64(
cutlass.Int64(_TMA_CACHE_EVICT_NORMAL).ir_value()
)
if cutlass.const_expr(self.enable_cache_hints):
if n_cluster_tiles >= m_cluster_tiles:
tma_cache_policy_a_sfa = cutlass.Int64(
cutlass.Int64(_TMA_CACHE_EVICT_LAST).ir_value()
)
tma_cache_policy_b_sfb = cutlass.Int64(
cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()
)
else:
tma_cache_policy_a_sfa = cutlass.Int64(
cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()
)
tma_cache_policy_b_sfb = cutlass.Int64(
cutlass.Int64(_TMA_CACHE_EVICT_LAST).ir_value()
)
#
# 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
)
# 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=tensormap_manager.get_tensormap_ptr(
tensormap_a_gmem_ptr,
cute.AddressSpace.generic,
),
cache_policy=tma_cache_policy_a_sfa,
)
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=tensormap_manager.get_tensormap_ptr(
tensormap_b_gmem_ptr,
cute.AddressSpace.generic,
),
cache_policy=tma_cache_policy_b_sfb,
)
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=tensormap_manager.get_tensormap_ptr(
tensormap_sfa_gmem_ptr,
cute.AddressSpace.generic,
),
cache_policy=tma_cache_policy_a_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=tensormap_manager.get_tensormap_ptr(
tensormap_sfb_gmem_ptr,
cute.AddressSpace.generic,
),
cache_policy=tma_cache_policy_b_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
)
last_group_idx = cur_group_idx
#
# Advance to next tile
#
tile_sched.advance_to_next_work()
work_tile = tile_sched.get_current_work()
#
# Wait A/B buffer empty
#
ab_pipeline.producer_tail(ab_producer_state)
#
# Specialized MMA warp
#
if warp_idx == self.mma_warp_id and initial_work_tile_info.is_valid_tile:
#
# Initialize tensormaps for A, B, SFA and SFB
#
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=self.sf_dtype,
)
# (MMA, MMA_M, MMA_K)
tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
tiled_mma,
self.mma_tiler,
self.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=self.sf_dtype,
)
# (MMA, MMA_N, MMA_K)
tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
tiled_mma,
self.mma_tiler,
self.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
#
work_tile = 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:
if cutlass.const_expr(
not (
self.problem_shape_n_const is None
or self.problem_shape_k_const is None
)
):
cur_k_tile_cnt = cluster_tile_count_k_const
else:
# MMA warp is only interested in number of tiles along K dimension
cur_k_tile_cnt, _ = group_gemm_ts_helper.search_cluster_tile_count_k(
work_tile.tile_idx, problem_shape_mnkl_smem
)
if cutlass.const_expr(self.overlapping_accum):
acc_stage_index = acc_producer_state.phase ^ 1
else:
acc_stage_index = acc_producer_state.index
# (MMA, MMA_M, MMA_N)
tCtAcc = tCtAcc_base[(None, None, None, acc_stage_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 on tCtAcc after first kblock
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 and initial_work_tile_info.is_valid_tile:
#
# Alloc tensor memory buffer
#
if warp_idx == self.epilog_warp_id[0]:
cute.arch.alloc_tmem(
self.num_tmem_alloc_cols,
tmem_holding_buf,
is_two_cta=use_2cta_instrs,
)
self.tmem_alloc_barrier.arrive_and_wait()
acc_tmem_ptr = cute.arch.retrieve_tmem_ptr(
self.acc_dtype,
alignment=16,
ptr_to_buffer_holding_addr=tmem_holding_buf,
)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
stg_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype)
work_tile = initial_work_tile_info
acc_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
)
while work_tile.is_valid_tile:
found_group = cutlass.Boolean(True)
cur_group_idx = cutlass.Int32(0)
cta_tile_idx_m = cutlass.Int32(0)
cta_tile_idx_n = cutlass.Int32(0)
cur_k_tile_cnt = cutlass.Int32(0)
problem_shape_m_group = cutlass.Int32(0)
problem_shape_n_group = cutlass.Int32(0)
problem_shape_k_group = cutlass.Int32(0)
if cutlass.const_expr(
not (
self.problem_shape_n_const is None
or self.problem_shape_k_const is None
)
):
(
found_group,
cur_group_idx,
cta_tile_idx_m,
cta_tile_idx_n,
problem_shape_m_group,
) = self.lookup_group_tile_info_dynamic_m(
work_tile.tile_idx,
group_count,
problem_shape_m_smem,
cluster_tile_count_mn_smem,
cluster_tile_count_n_const,
)
cur_k_tile_cnt = cluster_tile_count_k_const
problem_shape_n_group = problem_shape_n_const
problem_shape_k_group = problem_shape_k_const
else:
grouped_gemm_cta_tile_info = group_gemm_ts_helper.delinearize_z(
work_tile.tile_idx, problem_shape_mnkl_smem
)
cur_group_idx = grouped_gemm_cta_tile_info.group_idx
cta_tile_idx_m = grouped_gemm_cta_tile_info.cta_tile_idx_m
cta_tile_idx_n = grouped_gemm_cta_tile_info.cta_tile_idx_n
cur_k_tile_cnt = grouped_gemm_cta_tile_info.cta_tile_count_k
problem_shape_m_group = grouped_gemm_cta_tile_info.problem_shape_m
problem_shape_n_group = grouped_gemm_cta_tile_info.problem_shape_n
problem_shape_k_group = grouped_gemm_cta_tile_info.problem_shape_k
if found_group:
copy_atom_t2r = sm100_utils.get_tmem_load_op(
self.cta_tile_shape_mnk,
self.c_layout,
self.c_dtype,
self.acc_dtype,
epi_tile,
use_2cta_instrs,
)
tCtAcc_epi = cute.flat_divide(
tCtAcc_base[((None, None), 0, 0, None)],
epi_tile,
)
tiled_copy_t2r = tcgen05.make_tmem_copy(
copy_atom_t2r, tCtAcc_epi[(None, None, 0, 0, 0)]
)
thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
tTR_tAcc_base = thr_copy_t2r.partition_S(tCtAcc_epi)
real_tensor_c = self.make_tensor_abc_for_tensormap_update(
cur_group_idx,
self.c_dtype,
(
problem_shape_m_group,
problem_shape_n_group,
problem_shape_k_group,
),
strides_abc,
ptrs_abc,
2,
)
gC_mnl_group = cute.local_tile(
real_tensor_c,
cute.slice_(self.mma_tiler, (None, None, 0)),
(None, None, None),
)
tCgC_group = thr_mma.partition_C(gC_mnl_group)
tCgC_group_epi = cute.flat_divide(
tCgC_group[((None, None), 0, 0, None, None, None)],
epi_tile,
)
tTR_gC = thr_copy_t2r.partition_D(tCgC_group_epi)
# Predicates for SIMT/STG path using thread mapping offsets.
thr_mapping = cute.make_identity_tensor(
(self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1])
)
thr_mapping_mn = cute.flat_divide(thr_mapping, epi_tile)
m_thr_offset = thr_copy_t2r.partition_D(thr_mapping_mn)
m_thr_offset = cute.group_modes(
m_thr_offset, 3, cute.rank(m_thr_offset)
)
mma_tile_coord_mnl = (
cta_tile_idx_m // cute.size(tiled_mma.thr_id.shape),
cta_tile_idx_n,
0,
)
tTR_gC_tile = tTR_gC[
(None, None, None, None, None, *mma_tile_coord_mnl)
]
tTR_rAcc = cute.make_rmem_tensor(
tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape,
self.acc_dtype,
)
tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype)
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
tTR_tAcc = tTR_tAcc_base[
(None, None, None, None, None, acc_stage_index)
]
is_k_tile_cnt_zero = cur_k_tile_cnt == 0
if not is_k_tile_cnt_zero:
acc_pipeline.consumer_wait(acc_consumer_state)
tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
tTR_gC_tile = cute.group_modes(
tTR_gC_tile, 3, cute.rank(tTR_gC_tile)
)
tile_m_begin = cta_tile_idx_m * self.cta_tile_shape_mnk[0]
tile_n_begin = cta_tile_idx_n * self.cta_tile_shape_mnk[1]
is_full_cta_tile = (
(tile_m_begin + self.cta_tile_shape_mnk[0])
<= problem_shape_m_group
) and (
(tile_n_begin + self.cta_tile_shape_mnk[1])
<= problem_shape_n_group
)
subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
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
)
tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)]
if not is_k_tile_cnt_zero:
cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
tTR_rC.store(tTR_rAcc.load().to(self.c_dtype))
else:
tTR_rC.fill(0)
if cutlass.const_expr(self.overlapping_accum):
if (
(not is_k_tile_cnt_zero)
and subtile_idx == self.iter_acc_early_release_in_epilogue
):
# Ensure all pending TMEM reads are visible before reusing buffer.
cute.arch.fence_view_async_tmem_load()
with cute.arch.elect_one():
acc_pipeline.consumer_release(acc_consumer_state)
acc_consumer_state.advance()
tTR_gC_subtile = tTR_gC_tile[(None, None, None, real_subtile_idx)]
if is_full_cta_tile:
# Fast path for aligned in-bounds tiles: enables vectorized STG.
cute.copy(stg_atom, tTR_rC, tTR_gC_subtile)
else:
# Edge path: predicated store for OOB safety.
m_thr_slice = m_thr_offset[(None, None, None, real_subtile_idx)]
tTR_rC_flat = cute.flatten(tTR_rC)
tTR_gC_subtile_flat = cute.flatten(tTR_gC_subtile)
tTR_pC = cute.make_rmem_tensor(
cute.make_layout(tTR_rC_flat.shape),
cutlass.Boolean,
)
is_full_subtile = cutlass.Boolean(True)
is_any_subtile = cutlass.Boolean(False)
for i in cutlass.range(cute.size(tTR_pC), unroll_full=True):
in_bounds = (
m_thr_slice[(i)][0] + tile_m_begin
< problem_shape_m_group
) and (
m_thr_slice[(i)][1] + tile_n_begin
< problem_shape_n_group
)
tTR_pC[i] = in_bounds
is_full_subtile = is_full_subtile and in_bounds
is_any_subtile = is_any_subtile or in_bounds
if is_full_subtile:
cute.copy(stg_atom, tTR_rC, tTR_gC_subtile)
elif is_any_subtile:
cute.copy(
stg_atom,
tTR_rC_flat,
tTR_gC_subtile_flat,
pred=cute.flatten(tTR_pC),
)
else:
# Entire subtile is out-of-bounds; skip store.
pass
if cutlass.const_expr(not self.overlapping_accum):
if not is_k_tile_cnt_zero:
with cute.arch.elect_one():
acc_pipeline.consumer_release(acc_consumer_state)
acc_consumer_state.advance()
tile_sched.advance_to_next_work()
work_tile = tile_sched.get_current_work()
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 cutlass.const_expr(use_2cta_instrs):
cute.arch.mbarrier_arrive(
tmem_dealloc_mbar_ptr, cta_rank_in_cluster ^ 1
)
cute.arch.mbarrier_wait(tmem_dealloc_mbar_ptr, 0)
cute.arch.dealloc_tmem(
acc_tmem_ptr, self.num_tmem_alloc_cols, is_two_cta=use_2cta_instrs
)
@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)}"
)
ptr_align = 32 if cutlass.const_expr(tensor_index == 2) else 16
tensor_gmem_ptr = cute.make_ptr(
dtype, ptr_i64, cute.AddressSpace.gmem, assumed_align=ptr_align
)
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
# For STG.256, require 32B pointer alignment and aligned row-major strides.
m = problem_shape_mnk[0]
n = cute.assume(problem_shape_mnk[1], 16)
return cute.make_tensor(
tensor_gmem_ptr,
cute.make_layout(
(m, n, c1),
stride=(cute.assume(stride_mn, 16), 1, 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), self.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), self.sf_vec_size
)
return cute.make_tensor(
tensor_gmem_ptr,
sfb_layout,
)
@cute.jit
def lookup_group_tile_info_dynamic_m(
self,
cta_tile_coord: tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32],
group_count: cutlass.Constexpr[int],
problem_shape_m: cute.Tensor,
cluster_tile_count_mn: cute.Tensor,
cluster_tile_count_n_const: cutlass.Int32,
) -> tuple[
cutlass.Boolean,
cutlass.Int32,
cutlass.Int32,
cutlass.Int32,
cutlass.Int32,
]:
"""
Lightweight grouped lookup for dynamic-M/static-NK scheduling.
Returns (found, group_idx, cta_tile_idx_m, cta_tile_idx_n, problem_shape_m).
"""
lane_idx = cute.arch.lane_idx()
found = cutlass.Boolean(False)
group_idx = cutlass.Int32(0)
tile_count_prev_group = cutlass.Int32(0)
if lane_idx == 0:
tile_count_searched = cutlass.Int32(0)
for cur_group_idx in cutlass.range(group_count):
cur_cluster_tile_count = cluster_tile_count_mn[(cur_group_idx,)]
next_tile_count = tile_count_searched + cur_cluster_tile_count
if (not found) and (cta_tile_coord[2] < next_tile_count):
found = cutlass.Boolean(True)
group_idx = cur_group_idx
tile_count_prev_group = tile_count_searched
tile_count_searched = next_tile_count
found = cute.arch.shuffle_sync(found, 0)
group_idx = cute.arch.shuffle_sync(group_idx, 0)
tile_count_prev_group = cute.arch.shuffle_sync(tile_count_prev_group, 0)
cluster_tile_idx_in_current_group = cta_tile_coord[2] - tile_count_prev_group
cluster_tile_idx_m = cluster_tile_idx_in_current_group // cluster_tile_count_n_const
cluster_tile_idx_n = (
cluster_tile_idx_in_current_group
- cluster_tile_idx_m * cluster_tile_count_n_const
)
cta_tile_idx_m = (
cluster_tile_idx_m * self.cluster_shape_mn[0] + cta_tile_coord[0]
)
cta_tile_idx_n = (
cluster_tile_idx_n * self.cluster_shape_mn[1] + cta_tile_coord[1]
)
problem_shape_m_group = problem_shape_m[(group_idx,)]
return found, group_idx, cta_tile_idx_m, cta_tile_idx_n, problem_shape_m_group
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),
self.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,
self.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, self.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
@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]:
"""Computes the number of stages for ACC and A/B 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)
:rtype: tuple[int, int]
"""
# ACC stages
num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2
# Calculate smem layout and size for one stage of A, B, SFA, SFB
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
)
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
num_ab_stage = (smem_capacity - mbar_helpers_bytes) // ab_bytes_per_stage
return num_acc_stage, num_ab_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),
)
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
# Size of smem we reserved for mbarrier, tensor memory management and tensormap update
reserved_smem_bytes = 1024
bytes_per_tensormap = 128
num_tensormaps = 5
# size of smem used for tensor memory management
tensor_memory_management_bytes = 12
def _build_ptr_tensors_pinned(
abc_tensors: List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
sfasfb_reordered_tensors: List[Tuple[torch.Tensor, torch.Tensor]],
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Build pointer tensors via pinned CPU staging for fast H2D copy.
Returns:
tensor_of_abc_ptrs: torch tensor (num_groups, 3) containing data pointers
tensor_of_sf_ptrs: torch tensor (num_groups, 2) containing data pointers
"""
num_groups = len(abc_tensors)
# reuse pinned pointer staging buffers to cut host allocation overhead
abc_ptrs_cpu, sf_ptrs_cpu = _get_ptr_staging_buffers(num_groups)
for i in range(num_groups):
a_t, b_t, c_t = abc_tensors[i]
sfa_t, sfb_t = sfasfb_reordered_tensors[i]
abc_ptrs_cpu[i, 0] = a_t.data_ptr()
abc_ptrs_cpu[i, 1] = b_t.data_ptr()
abc_ptrs_cpu[i, 2] = c_t.data_ptr()
sf_ptrs_cpu[i, 0] = sfa_t.data_ptr()
sf_ptrs_cpu[i, 1] = sfb_t.data_ptr()
tensor_of_abc_ptrs = abc_ptrs_cpu.to(device="cuda", non_blocking=True)
tensor_of_sf_ptrs = sf_ptrs_cpu.to(device="cuda", non_blocking=True)
return tensor_of_abc_ptrs, tensor_of_sf_ptrs
def _build_stride_tensor_pinned(
problem_sizes: List[Tuple[int, int, int, int]],
device: str,
) -> torch.Tensor:
"""
Create torch tensor containing stride data using pinned CPU staging.
Returns:
tensor_of_strides_abc: torch tensor (num_groups, 3, 2) containing strides
"""
ps = torch.tensor(problem_sizes, dtype=torch.int32, device="cpu")
n = ps[:, 1]
k = ps[:, 2]
num_groups = ps.shape[0]
strides_cpu = torch.empty(
(num_groups, 3, 2), dtype=torch.int32, pin_memory=True
)
strides_cpu[:, 0, 0] = k
strides_cpu[:, 0, 1] = 1
strides_cpu[:, 1, 0] = k
strides_cpu[:, 1, 1] = 1
strides_cpu[:, 2, 0] = n
strides_cpu[:, 2, 1] = 1
return strides_cpu.to(device=device, non_blocking=True)
@functools.lru_cache(maxsize=32)
def _get_stride_tensor_cached(
problem_sizes: tuple[tuple[int, int, int, int], ...],
device: str,
) -> torch.Tensor:
return _build_stride_tensor_pinned(list(problem_sizes), device)
def _build_problem_shape_tensor_pinned(
problem_sizes: List[Tuple[int, int, int, int]],
device: str,
) -> torch.Tensor:
"""
Create torch tensor containing problem shapes using pinned CPU staging.
Returns:
tensor_of_problem_shapes: torch tensor (num_groups, 4) containing (m, n, k, l)
"""
problem_shapes_cpu = torch.tensor(problem_sizes, dtype=torch.int32, pin_memory=True)
return problem_shapes_cpu.to(device=device, non_blocking=True)
@functools.lru_cache(maxsize=32)
def _get_problem_shape_tensor_cached(
problem_sizes: tuple[tuple[int, int, int, int], ...],
device: str,
) -> torch.Tensor:
return _build_problem_shape_tensor_pinned(list(problem_sizes), device)
def _build_group_runtime_meta_tensor_pinned(
problem_sizes: List[Tuple[int, int, int, int]],
cluster_tile_shape_mn: Tuple[int, int],
device: str,
) -> torch.Tensor:
"""
Build runtime metadata tensor (group_count, 4): [M, cluster_tile_count_mn, 0, 0].
"""
ps = torch.tensor(problem_sizes, dtype=torch.int32, pin_memory=True)
m_cpu = ps[:, 0].contiguous()
n_cpu = ps[:, 1]
cluster_count_m = (m_cpu + cluster_tile_shape_mn[0] - 1) // cluster_tile_shape_mn[0]
cluster_count_n = (n_cpu + cluster_tile_shape_mn[1] - 1) // cluster_tile_shape_mn[1]
cluster_tile_count_mn_cpu = (cluster_count_m * cluster_count_n).contiguous()
runtime_meta_cpu = torch.zeros((ps.shape[0], 4), dtype=torch.int32, pin_memory=True)
runtime_meta_cpu[:, 0] = m_cpu
runtime_meta_cpu[:, 1] = cluster_tile_count_mn_cpu
return runtime_meta_cpu.to(device=device, non_blocking=True)
@functools.lru_cache(maxsize=32)
def _get_group_runtime_meta_tensor_cached(
problem_sizes: tuple[tuple[int, int, int, int], ...],
cluster_tile_shape_mn: tuple[int, int],
device: str,
) -> torch.Tensor:
return _build_group_runtime_meta_tensor_pinned(
list(problem_sizes), cluster_tile_shape_mn, device
)
_tensormap_cache: dict[str, torch.Tensor] = {}
_ptr_staging_cache: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
_last_compile_key: Tuple[int, tuple[tuple[int, int, int, int], ...]] | None = None
def _select_ab_stage_cap_for_family(
num_groups: int,
problem_sizes: tuple[tuple[int, int, int, int], ...],
) -> int | None:
"""
Select optional AB-stage cap based on static (g, n, k) family.
"""
n_values = {n for _, n, _, _ in problem_sizes}
k_values = {k for _, _, k, _ in problem_sizes}
if len(n_values) != 1 or len(k_values) != 1:
return None
n = next(iter(n_values))
k = next(iter(k_values))
# Benchmark families:
# 1) g=8, n=4096, k=7168 -> cap=5
# 2) g=8, n=7168, k=2048 -> cap=4
# 3) g=2, n=3072, k=4096 -> cap=4
# 4) g=2, n=4096, k=1536 -> cap=4
if num_groups == 8 and n == 4096 and k == 7168:
return None # 7+ fails best 6
if num_groups == 8 and n == 7168 and k == 2048:
return None # 5+ fails best 4
if num_groups == 2 and n == 3072 and k == 4096:
return None # fails best 3
if num_groups == 2 and n == 4096 and k == 1536:
return None # fails best 3
return None
def _try_resolve_benchmark_family(
num_groups: int,
problem_sizes: tuple[tuple[int, int, int, int], ...],
) -> tuple[int, int, int] | None:
if len(problem_sizes) == 0:
raise ValueError("problem_sizes must be non-empty")
n0 = problem_sizes[0][1]
k0 = problem_sizes[0][2]
for _, n, k, _ in problem_sizes:
if n != n0 or k != k0:
return None
family_key = (num_groups, n0, k0)
family_id = _BENCHMARK_FAMILY_BY_KEY.get(family_key)
if family_id is None:
return None
return family_id, n0, k0
def get_tensormap_tensor(device: str = "cuda") -> torch.Tensor:
"""
Create tensormap buffer.
Returns:
tensor_of_tensormap: torch tensor (148, 5, 16) for tensormap storage
"""
if device in _tensormap_cache:
return _tensormap_cache[device]
sm_count = 148
tensormap_shape = (
sm_count,
Sm100GroupedBlockScaledGemmKernel.num_tensormaps,
Sm100GroupedBlockScaledGemmKernel.bytes_per_tensormap // 8,
)
tensor_of_tensormap = torch.empty(
tensormap_shape, dtype=torch.int64, device=device
)
_tensormap_cache[device] = tensor_of_tensormap
return tensor_of_tensormap
def _get_ptr_staging_buffers(
num_groups: int,
) -> tuple[torch.Tensor, torch.Tensor]:
staging = _ptr_staging_cache.get(num_groups)
if staging is None:
abc_ptrs_cpu = torch.empty((num_groups, 3), dtype=torch.int64, pin_memory=True)
sf_ptrs_cpu = torch.empty((num_groups, 2), dtype=torch.int64, pin_memory=True)
staging = (abc_ptrs_cpu, sf_ptrs_cpu)
_ptr_staging_cache[num_groups] = staging
return staging
@functools.lru_cache(maxsize=32)
def compile_kernel(
num_groups: int,
problem_sizes: tuple[tuple[int, int, int, int], ...],
):
"""
Compile the kernel once and cache it using compile-time parameters as the key.
This should be called before any timing measurements.
Args:
num_groups: Number of GEMM groups
problem_sizes: Tuple of (m, n, k, l) tuples for each group (must be hashable)
Returns:
The compiled TVM FFI kernel function
"""
family = _try_resolve_benchmark_family(num_groups, problem_sizes)
if family is None:
all_n_7168 = all(n == 7168 for _, n, _, _ in problem_sizes)
mma_tiler_mn = (128, 256) if all_n_7168 else (128, 128)
problem_shape_n_const = None
problem_shape_k_const = None
enable_cache_hints = (
num_groups == 8
and all(n == 4096 for _, n, _, _ in problem_sizes)
and all(k == 7168 for _, _, k, _ in problem_sizes)
)
else:
family_id, _, _ = family
(
_,
problem_shape_n_const,
problem_shape_k_const,
mma_tiler_mn,
enable_cache_hints,
) = _BENCHMARK_FAMILY_SPECS[family_id]
cluster_shape_mn = (1, 2)
use_2cta_instrs = mma_tiler_mn[0] == 256
sm_count = 148
max_active_clusters = 148 // prod(cluster_shape_mn)
num_tensormap_buffers = sm_count
# Compute total number of clusters needed across all groups
cta_tile_shape_mn = [mma_tiler_mn[0], mma_tiler_mn[1]]
cluster_tile_shape_mn = tuple(
x * y for x, y in zip(cta_tile_shape_mn, cluster_shape_mn)
)
total_num_clusters = 0
for m, n, _, _ in problem_sizes:
num_clusters_mn = tuple(
(x + y - 1) // y for x, y in zip((m, n), cluster_tile_shape_mn)
)
total_num_clusters += functools.reduce(lambda x, y: x * y, num_clusters_mn)
ab_stage_cap = _select_ab_stage_cap_for_family(num_groups, problem_sizes)
# Create kernel instance
grouped_blockscaled_gemm = Sm100GroupedBlockScaledGemmKernel(
sf_vec_size,
mma_tiler_mn,
cluster_shape_mn,
enable_cache_hints,
ab_stage_cap=ab_stage_cap,
problem_shape_n_const=problem_shape_n_const,
problem_shape_k_const=problem_shape_k_const,
)
# Create fake tensors for TVM FFI compilation (row-major layout)
fake_problem_shape = cute.runtime.make_fake_compact_tensor(
cutlass.Int32, (num_groups, 4),
stride_order=(1, 0), assumed_align=16,
)
fake_group_runtime_meta = cute.runtime.make_fake_compact_tensor(
cutlass.Int32, (num_groups, 4),
stride_order=(1, 0), assumed_align=16,
)
fake_strides = cute.runtime.make_fake_compact_tensor(
cutlass.Int32, (num_groups, 3, 2),
stride_order=(2, 1, 0), assumed_align=16,
)
fake_abc_addr = cute.runtime.make_fake_compact_tensor(
cutlass.Int64, (num_groups, 3),
stride_order=(1, 0), assumed_align=16,
)
fake_sf_addr = cute.runtime.make_fake_compact_tensor(
cutlass.Int64, (num_groups, 2),
stride_order=(1, 0), assumed_align=16,
)
fake_tensormap = cute.runtime.make_fake_compact_tensor(
cutlass.Int64, (num_tensormap_buffers, 5, 16),
stride_order=(2, 1, 0), assumed_align=16,
)
# Compile grouped GEMM kernel with TVM FFI
compiled_func = cute.compile(
grouped_blockscaled_gemm,
num_groups, # group_count: Constexpr
fake_problem_shape, # problem_shape_mnkl: Tensor
fake_group_runtime_meta, # group_runtime_meta: Tensor
use_2cta_instrs, # use_2cta_instrs: Constexpr
fake_strides, # strides_abc: Tensor
fake_abc_addr, # tensor_address_abc: Tensor
fake_sf_addr, # tensor_address_sfasfb: Tensor
total_num_clusters, # total_num_clusters: Constexpr
fake_tensormap, # tensormap_cute_tensor: Tensor
max_active_clusters, # max_active_clusters: Constexpr
options="--enable-tvm-ffi --opt-level 2 --generate-line-info",
)
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 prepares torch tensors, compiles with TVM FFI, and launches the kernel.
Args:
data: Tuple of (abc_tensors, sfasfb_tensors, sfasfb_reordered_tensors, problem_sizes) where:
abc_tensors: list of tuples (a, b, c) for each group
sfasfb_tensors: list of tuples (sfa, sfb) - original scale factors
sfasfb_reordered_tensors: list of tuples (sfa_reordered, sfb_reordered) - reordered scale factors
problem_sizes: list of tuples (m, n, k, l)
Returns:
list of c tensors where c is torch.Tensor[float16] of shape [m, n, l] for each group
"""
abc_tensors: List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
problem_sizes: List[Tuple[int, int, int, int]] = []
global _last_compile_key
abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data
num_groups = len(problem_sizes)
tuple_of_problem_sizes = tuple(tuple(ps) for ps in problem_sizes)
compile_key = (num_groups, tuple_of_problem_sizes)
# Prepare pointers using pinned CPU staging for fast H2D copy
torch_tensor_abc, torch_tensor_sf = _build_ptr_tensors_pinned(
abc_tensors, sfasfb_reordered_tensors
)
# Get stride tensor from cache (keyed by problem sizes)
torch_tensor_strides = _get_stride_tensor_cached(
tuple_of_problem_sizes,
"cuda",
)
torch_tensor_problem_sizes = _get_problem_shape_tensor_cached(
tuple_of_problem_sizes,
"cuda",
)
all_n_7168 = all(n == 7168 for _, n, _, _ in tuple_of_problem_sizes)
mma_tiler_mn = (128, 256) if all_n_7168 else (128, 128)
cluster_shape_mn = (1, 2)
cluster_tile_shape_mn = tuple(
x * y for x, y in zip(mma_tiler_mn, cluster_shape_mn)
)
torch_tensor_group_runtime_meta = _get_group_runtime_meta_tensor_cached(
tuple_of_problem_sizes,
cluster_tile_shape_mn,
"cuda",
)
# Get tensormap tensor for runtime
torch_tensor_tensormap = get_tensormap_tensor("cuda")
# Get compiled kernel from cache
if _last_compile_key != compile_key:
compile_kernel.cache_clear()
_last_compile_key = compile_key
compiled_grouped_gemm = compile_kernel(
num_groups,
tuple_of_problem_sizes,
)
# Call the compiled TVM FFI kernel with torch tensors directly
compiled_grouped_gemm(
torch_tensor_problem_sizes, # problem_shape_mnkl tensor
torch_tensor_group_runtime_meta, # group_runtime_meta tensor
torch_tensor_strides, # strides_abc tensor
torch_tensor_abc, # tensor_address_abc tensor
torch_tensor_sf, # tensor_address_sfasfb tensor
torch_tensor_tensormap, # tensormap tensor
)
# Return the C tensors
res = []
for i in range(num_groups):
res.append(abc_tensors[i][2])
return res
scrolls · 2787 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 500515.
⋯ diff truncated: revisions differ almost entirely
Best evidence level for this revision: reported
JSON