submission 384707
yue · python · License unknown
Use it
Vendorable · source mirrored · license unknownView source →
No package. Vendor the mirrored source: 2659 lines, June 9 Researcher Reciprocity License v1.0.
v2_test.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-modal-nvfp4-dual-gemm-384707?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:90ebd339377e64acf475bbed365acbe9557e5a4a5997ec8d33848318aa31473a
license declaredunknown
license concludedunknown
authorsyue
imported2026-08-15
Techniques
Extracted from the mirrored source by pattern, never inferred. Each row cites its line.
fp4
tiled_mma, # Tiled MMA object defining NVFP4 GEMM compute patternfused-epilogue
self.epi_tile = sm100_utils.compute_epilogue_tile_shape(mbarrier
self.epilog_sync_barrier = pipeline.NamedBarrier(persistent-kernel
GPU device kernel performing batched GEMM computation (non-persistent).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
v2_test.py2659 lines
import argparse
from typing import Type, Tuple, Union, Optional
import cuda.bindings.driver as cuda
import torch
import math
import cutlass
import cutlass.cute as cute
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.pipeline as pipeline
import cutlass.utils.blackwell_helpers as sm100_utils
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.cute.runtime import from_dlpack
from torch._higher_order_ops.torchbind import call_torchbind_fake
import cuda.bindings.driver as cuda
from cutlass import Float16, Float32, Int8, Int16, Int32, const_expr
from cutlass._mlir import ir
from cutlass._mlir.dialects import nvvm
from cutlass.cutlass_dsl import T
from task import input_t, output_t
from cutlass.cute.runtime import make_ptr
# Constant for silu: -log2(e) for converting exp to exp2
NEG_LOG2_E = -math.log2(math.e)
# Shape of the K dimension for the MMA instruction
mma_inst_shape_k = 64
# 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
# Number of threads per CUDA thread block
threads_per_cta = 128
# Total number of columns in tmem
num_tmem_alloc_cols = 512
class DualGemmKernel:
def __init__(
self,
sf_vec_size: int,
mma_tiler_mn: Tuple[int, int],
cluster_shape_mn: Tuple[int, int],
):
self.ab_dtype = cutlass.Float4E2M1FN
self.sf_dtype = cutlass.Float8E4M3FN
self.c_dtype = cutlass.Float16
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
# 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.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 id 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)),
)
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
def _setup_attributes(self):
# 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.cta_tile_shape_mnk_sfb = (
self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape),
self.mma_tiler_sfb[1],
self.mma_tiler_sfb[2],
)
# Compute cluster layout
self.cluster_layout_vmnk = cute.tiled_divide(
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,
)
self.epi_tile_n = cute.size(self.epi_tile[1])
self.num_acc_stage, self.num_ab_stage, self.num_c_stage = (1, 5, 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,
)
self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(
self.c_dtype,
self.c_layout,
self.epi_tile,
self.num_c_stage,
)
# Compute number of TMEM columns for SFA/SFB/Accumulator
sf_atom_mn = 32
self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k
self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k
self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols
self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] * self.num_acc_stage
def compute_silu_mul(self, x0, x1, y0, y1):
"""
Compute silu(x) * y for two float pairs.
silu(x) = x * sigmoid(x) = x / (1 + exp(-x))
Pipeline:
1. exp_arg = x * NEG_LOG2_E (convert to exp2 domain)
2. exp_val = exp2(exp_arg) (fast exp2 approximation)
3. denom = exp_val + 1.0 (denominator of sigmoid)
4. sigmoid = rcp(denom) (fast reciprocal)
5. result = sigmoid * y * x (silu(x) * y)
Returns: (result0, result1)
"""
# Stage 1: Multiply by -log2(e) to convert exp(-x) to exp2
t0, t1 = cute.arch.mul_packed_f32x2(
(x0, x1), (NEG_LOG2_E, NEG_LOG2_E), rnd=nvvm.RoundingModeKind.RN
)
# Stage 2: Fast exp2 approximation
t0 = cute.math.exp2(t0, fastmath=True)
t1 = cute.math.exp2(t1, fastmath=True)
# Stage 3: Add 1.0 to get denominator
t0, t1 = cute.arch.add_packed_f32x2(
(t0, t1), (1.0, 1.0), rnd=nvvm.RoundingModeKind.RN
)
# Stage 4: Fast reciprocal for sigmoid
t0 = cute.arch.rcp_approx(t0)
t1 = cute.arch.rcp_approx(t1)
# Stage 5: Multiply by y
t0, t1 = cute.arch.mul_packed_f32x2(
(t0, t1), (y0, y1), rnd=nvvm.RoundingModeKind.RN
)
# Stage 6: Multiply by x for final result
t0, t1 = cute.arch.mul_packed_f32x2(
(t0, t1), (x0, x1), rnd=nvvm.RoundingModeKind.RN
)
return t0, t1
@cute.jit
def __call__(
self,
a_ptr: cute.Pointer,
b1_ptr: cute.Pointer,
b2_ptr: cute.Pointer,
sfa_ptr: cute.Pointer,
sfb1_ptr: cute.Pointer,
sfb2_ptr: cute.Pointer,
c_ptr: cute.Pointer,
m: cutlass.Constexpr,
n: cutlass.Constexpr,
k: cutlass.Constexpr,
l: cutlass.Constexpr,
):
# Setup static attributes before smem/grid/tma computation
self.a_dtype: Type[cutlass.Numeric] = ab_dtype
self.b_dtype: Type[cutlass.Numeric] = ab_dtype
self.sf_dtype: Type[cutlass.Numeric] = sf_dtype
self.c_dtype: Type[cutlass.Numeric] = c_dtype
self.a_major_mode = tcgen05.OperandMajorMode.K
self.b_major_mode = tcgen05.OperandMajorMode.K
self.c_layout = utils.LayoutEnum.ROW_MAJOR
# Setup attributes that dependent on gemm inputs
self._setup_attributes()
a_tensor = cute.make_tensor(
a_ptr,
cute.make_ordered_layout((m, k, l), order=(1, 0, 2)),
)
b1_tensor = cute.make_tensor(
b1_ptr,
cute.make_ordered_layout((n, k, l), order=(1, 0, 2)),
)
b2_tensor = cute.make_tensor(
b2_ptr,
cute.make_ordered_layout((n, k, l), order=(1, 0, 2)),
)
c_tensor = cute.make_tensor(
c_ptr, cute.make_ordered_layout((m, n, l), order=(1, 0, 2))
)
# 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(
a_tensor.shape, self.sf_vec_size
)
sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
# ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
b1_tensor.shape, self.sf_vec_size
)
sfb1_tensor = cute.make_tensor(sfb1_ptr, sfb_layout)
sfb2_tensor = cute.make_tensor(sfb2_ptr, 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,
a_tensor,
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_b1, tma_tensor_b1 = cute.nvgpu.make_tiled_tma_atom_B(
b_op,
b1_tensor,
b_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
)
tma_atom_b2, tma_tensor_b2 = cute.nvgpu.make_tiled_tma_atom_B(
b_op,
b2_tensor,
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,
sfa_tensor,
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_sfb1, tma_tensor_sfb1 = cute.nvgpu.make_tiled_tma_atom_B(
sfb_op,
sfb1_tensor,
sfb_smem_layout,
self.mma_tiler_sfb,
tiled_mma_sfb,
self.cluster_layout_sfb_vmnk.shape,
internal_type=cutlass.Int16,
)
tma_atom_sfb2, tma_tensor_sfb2 = cute.nvgpu.make_tiled_tma_atom_B(
sfb_op,
sfb2_tensor,
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 * 2 + sfa_copy_size + sfb_copy_size * 2
) * 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(),
c_tensor,
epi_smem_layout,
self.epi_tile,
)
self.buffer_align_bytes = 1024
# Compute grid size
# grid_m is calculated by ceil_div(m, cta_tile_shape_mnk[0])
# For 2-CTA instructions, cta_tile_shape_mnk[0] = mma_tiler_mn[0] / 2
# grid_n is calculated by ceil_div(n, mma_tiler_mn[1])
# grid_l is calculated by l
grid_m = cute.ceil_div(m, self.cta_tile_shape_mnk[0])
grid_n = cute.ceil_div(n, self.mma_inst_shape_mn[1])
grid_l = l
grid = (grid_m, grid_n, grid_l)
# Define shared storage for kernel
@cute.struct
class SharedStorage:
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
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sC: cute.struct.Align[
cute.struct.MemRange[
self.c_dtype,
cute.cosize(self.c_smem_layout_staged.outer),
],
self.buffer_align_bytes,
]
# (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)
sB1: cute.struct.Align[
cute.struct.MemRange[
self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)
],
self.buffer_align_bytes,
]
sB2: 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)
sSFB1: cute.struct.Align[
cute.struct.MemRange[
self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)
],
self.buffer_align_bytes,
]
sSFB2: 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_b1,
tma_tensor_b1,
tma_atom_b2,
tma_tensor_b2,
tma_atom_sfa,
tma_tensor_sfa,
tma_atom_sfb1,
tma_tensor_sfb1,
tma_atom_sfb2,
tma_tensor_sfb2,
tma_atom_c,
tma_tensor_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.c_smem_layout_staged,
self.epi_tile,
).launch(
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
)
return
# GPU device kernel
@cute.kernel
def kernel(
self,
tiled_mma: cute.TiledMma,
tiled_mma_sfb: cute.TiledMma, # Receive SFB Tiled MMA
tma_atom_a: cute.CopyAtom,
mA_mkl: cute.Tensor,
tma_atom_b1: cute.CopyAtom,
mB1_nkl: cute.Tensor,
tma_atom_b2: cute.CopyAtom,
mB2_nkl: cute.Tensor,
tma_atom_sfa: cute.CopyAtom,
mSFA_mkl: cute.Tensor,
tma_atom_sfb1: cute.CopyAtom,
mSFB1_nkl: cute.Tensor,
tma_atom_sfb2: cute.CopyAtom,
mSFB2_nkl: cute.Tensor,
tma_atom_c: Optional[cute.CopyAtom],
mC_mnl: cute.Tensor,
cluster_layout_vmnk: cute.Layout,
cluster_layout_sfb_vmnk: cute.Layout, # Receive SFB Cluster 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, None],
epi_tile: cute.Tile,
):
"""
GPU device kernel performing batched GEMM computation (non-persistent).
"""
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
#
# Prefetch tma desc
#
if warp_idx == self.tma_warp_id:
cpasync.prefetch_descriptor(tma_atom_a)
cpasync.prefetch_descriptor(tma_atom_b1)
cpasync.prefetch_descriptor(tma_atom_b2)
cpasync.prefetch_descriptor(tma_atom_sfa)
cpasync.prefetch_descriptor(tma_atom_sfb1)
cpasync.prefetch_descriptor(tma_atom_sfb2)
cpasync.prefetch_descriptor(tma_atom_c)
use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2
#
# 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: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier
#
smem = utils.SmemAllocator()
storage = smem.allocate(self.shared_storage)
# 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,
defer_sync=True,
)
# 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,
defer_sync=True,
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
barrier_for_retrieve=self.tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
)
# 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
#
# (EPI_TILE_M, EPI_TILE_N, STAGE)
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)
sB1 = storage.sB1.get_tensor(
b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner
)
sB2 = storage.sB2.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)
sSFB1 = storage.sSFB1.get_tensor(sfb_smem_layout_staged)
sSFB2 = storage.sSFB2.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)
gB1_nkl = cute.local_tile(
mB1_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
# (bN, bK, RestN, RestK, RestL)
gB2_nkl = cute.local_tile(
mB2_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)
gSFB1_nkl = cute.local_tile(
mSFB1_nkl,
cute.slice_(self.mma_tiler_sfb, (0, None, None)),
(None, None, None),
)
# (bN, bK, RestN, RestK, RestL)
gSFB2_nkl = cute.local_tile(
mSFB2_nkl,
cute.slice_(self.mma_tiler_sfb, (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)
)
k_tile_cnt = cute.size(gA_mkl, mode=[3])
#
# 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)
tCgB1 = thr_mma.partition_B(gB1_nkl)
tCgB2 = thr_mma.partition_B(gB2_nkl)
# (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
tCgSFA = thr_mma.partition_A(gSFA_mkl)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgSFB1 = thr_mma_sfb.partition_B(gSFB1_nkl)
tCgSFB2 = thr_mma_sfb.partition_B(gSFB2_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)
tBsB1, tBgB1 = cpasync.tma_partition(
tma_atom_b1,
block_in_cluster_coord_vmnk[1],
b_cta_layout,
cute.group_modes(sB1, 0, 3),
cute.group_modes(tCgB1, 0, 3),
)
tBsB2, tBgB2 = cpasync.tma_partition(
tma_atom_b2,
block_in_cluster_coord_vmnk[1],
b_cta_layout,
cute.group_modes(sB2, 0, 3),
cute.group_modes(tCgB2, 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)
tBsSFB1, tBgSFB1 = cute.nvgpu.cpasync.tma_partition(
tma_atom_sfb1,
block_in_cluster_coord_sfb_vmnk[1],
sfb_cta_layout,
cute.group_modes(sSFB1, 0, 3),
cute.group_modes(tCgSFB1, 0, 3),
)
tBsSFB1 = cute.filter_zeros(tBsSFB1)
tBgSFB1 = cute.filter_zeros(tBgSFB1)
tBsSFB2, tBgSFB2 = cute.nvgpu.cpasync.tma_partition(
tma_atom_sfb2,
block_in_cluster_coord_sfb_vmnk[1],
sfb_cta_layout,
cute.group_modes(sSFB2, 0, 3),
cute.group_modes(tCgSFB2, 0, 3),
)
tBsSFB2 = cute.filter_zeros(tBsSFB2)
tBgSFB2 = cute.filter_zeros(tBgSFB2)
#
# 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)
tCrB1 = tiled_mma.make_fragment_B(sB1)
tCrB2 = tiled_mma.make_fragment_B(sB2)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
# (MMA, MMA_M, MMA_N) - no STAGE dimension, as cute.gemm expects rank 2 or 3
tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
#
# Specialized TMA load warp
#
if warp_idx == self.tma_warp_id:
#
# Get tile coord from block index (non-persistent: one tile per CTA)
#
block_idx = cute.arch.block_idx()
mma_tile_coord_mnl = (
block_idx[0] // cute.size(tiled_mma.thr_id.shape),
block_idx[1],
block_idx[2],
)
ab_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_ab_stage
)
#
# 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)
tBgB1_slice = tBgB1[
(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])
]
# ((atom_v, rest_v), RestK)
tBgB2_slice = tBgB2[
(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])
]
slice_n = mma_tile_coord_mnl[1]
if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
slice_n = mma_tile_coord_mnl[1] // 2
# ((atom_v, rest_v), RestK)
tBgSFB1_slice = tBgSFB1[
(None, slice_n, None, mma_tile_coord_mnl[2])
]
# ((atom_v, rest_v), RestK)
tBgSFB2_slice = tBgSFB2[
(None, slice_n, None, mma_tile_coord_mnl[2])
]
# Initial prefetch: prefetch first num_ab_stage tiles
for prefetch_tile in cutlass.range(0, self.num_ab_stage, unroll=1):
cute.prefetch(tma_atom_a, tAgA_slice[(None, prefetch_tile)])
cute.prefetch(tma_atom_b1, tBgB1_slice[(None, prefetch_tile)])
cute.prefetch(tma_atom_b2, tBgB2_slice[(None, prefetch_tile)])
cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, prefetch_tile)])
cute.prefetch(tma_atom_sfb1, tBgSFB1_slice[(None, prefetch_tile)])
cute.prefetch(tma_atom_sfb2, tBgSFB2_slice[(None, prefetch_tile)])
# Peek (try_wait) AB buffer empty for k_tile = 0
ab_producer_state.reset_count()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < k_tile_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(
ab_producer_state
)
#
# Tma load loop
#
for k_tile in cutlass.range(0, k_tile_cnt, 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,
)
cute.copy(
tma_atom_b1,
tBgB1_slice[(None, ab_producer_state.count)],
tBsB1[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=b_full_mcast_mask,
)
cute.copy(
tma_atom_b2,
tBgB2_slice[(None, ab_producer_state.count)],
tBsB2[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=b_full_mcast_mask,
)
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,
)
cute.copy(
tma_atom_sfb1,
tBgSFB1_slice[(None, ab_producer_state.count)],
tBsSFB1[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=sfb_full_mcast_mask,
)
cute.copy(
tma_atom_sfb2,
tBgSFB2_slice[(None, ab_producer_state.count)],
tBsSFB2[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=sfb_full_mcast_mask,
)
# Prefetch next tiles ahead by num_ab_stage
prefetch_idx = ab_producer_state.count + self.num_ab_stage
if prefetch_idx < k_tile_cnt:
cute.prefetch(tma_atom_a, tAgA_slice[(None, prefetch_idx)])
cute.prefetch(tma_atom_b1, tBgB1_slice[(None, prefetch_idx)])
cute.prefetch(tma_atom_b2, tBgB2_slice[(None, prefetch_idx)])
cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, prefetch_idx)])
cute.prefetch(tma_atom_sfb1, tBgSFB1_slice[(None, prefetch_idx)])
cute.prefetch(tma_atom_sfb2, tBgSFB2_slice[(None, prefetch_idx)])
# Peek (try_wait) AB buffer empty for next iteration
ab_producer_state.advance()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < k_tile_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(
ab_producer_state
)
#
# Wait A/B buffer empty
#
ab_pipeline.producer_tail(ab_producer_state)
# ---------- MMA warp: AB consumer + GEMM + ACC producer ----------
elif warp_idx == self.mma_warp_id:
#
# Bar sync for retrieve tensor memory ptr from shared mem
#
tmem.wait_for_alloc()
#
# Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor
#
acc1_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
# Make accumulator tmem tensor
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc1 = cute.make_tensor(acc1_tmem_ptr, tCtAcc_fake.layout)
acc2_tmem_ptr = cute.recast_ptr(
acc1_tmem_ptr + self.num_accumulator_tmem_cols,
dtype=self.acc_dtype,
)
tCtAcc2 = cute.make_tensor(acc2_tmem_ptr, tCtAcc_fake.layout)
# Make SFA tmem tensor
sfa_tmem_ptr = cute.recast_ptr(
acc1_tmem_ptr + self.num_accumulator_tmem_cols * 2,
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
# 16 is determined by BLOCK_K, BLOCK_K = 64, which is 4 scale factors,
# * 4 (4 tmem columns) = 16.
sfb1_tmem_ptr = cute.recast_ptr(
acc1_tmem_ptr + self.num_accumulator_tmem_cols * 2 + 16,
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)),
)
tCtSFB1 = cute.make_tensor(sfb1_tmem_ptr, tCtSFB_layout)
sfb2_tmem_ptr = cute.recast_ptr(
acc1_tmem_ptr + self.num_accumulator_tmem_cols * 2 + 16 * 2,
dtype=self.sf_dtype,
)
tCtSFB2 = cute.make_tensor(sfb2_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_sfb1,
tCsSFB1_compact_s2t,
tCtSFB1_compact_s2t,
) = self.mainloop_s2t_copy_and_partition(sSFB1, tCtSFB1)
(
tiled_copy_s2t_sfb2,
tCsSFB2_compact_s2t,
tCtSFB2_compact_s2t,
) = self.mainloop_s2t_copy_and_partition(sSFB2, tCtSFB2)
#
# Get tile coord from block index (non-persistent: one tile per CTA)
#
block_idx = cute.arch.block_idx()
mma_tile_coord_mnl = (
block_idx[0] // cute.size(tiled_mma.thr_id.shape),
block_idx[1],
block_idx[2],
)
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
)
# Get accumulator stage index
acc_stage_index = 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 < 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)
tCtSFB1_mma = tCtSFB1
tCtSFB2_mma = tCtSFB2
if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
# If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words (ignores first 64 columns of SFB)
offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0)
tCtSFB1_mma = cute.make_tensor(sfb1_tmem_ptr + offset, tCtSFB_layout)
tCtSFB2_mma = cute.make_tensor(sfb2_tmem_ptr + offset, tCtSFB_layout)
elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
# Move in increments of 64 columns of SFB
# when block_n = 64, we still load 128 rows (along N direction of SFB) of SFB,
# the tile index == 0 load to first 64 rows, the tile index == 1 load to second 64 rows.
# This offset is for that.
# And because the tmem cols layout here https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-a-layout-4x
# The col offset by 8 is the second 64 rows of SFB.
offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 8)
tCtSFB1_mma = cute.make_tensor(sfb1_tmem_ptr + offset, tCtSFB_layout)
tCtSFB2_mma = cute.make_tensor(sfb2_tmem_ptr + offset, tCtSFB_layout)
#
# Reset the ACCUMULATE field for each tile
#
tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
#
# Mma mainloop
#
for k_tile in range(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]
tCsSFB1_compact_s2t_staged = tCsSFB1_compact_s2t[s2t_stage_coord]
tCsSFB2_compact_s2t_staged = tCsSFB2_compact_s2t[s2t_stage_coord]
cute.copy(
tiled_copy_s2t_sfa,
tCsSFA_compact_s2t_staged,
tCtSFA_compact_s2t,
)
cute.copy(
tiled_copy_s2t_sfb1,
tCsSFB1_compact_s2t_staged,
tCtSFB1_compact_s2t,
)
cute.copy(
tiled_copy_s2t_sfb2,
tCsSFB2_compact_s2t_staged,
tCtSFB2_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,
tCtSFB1_mma[sf_kblock_coord].iterator,
)
cute.gemm(
tiled_mma,
tCtAcc1,
tCrA[kblock_coord],
tCrB1[kblock_coord],
tCtAcc1,
)
tiled_mma.set(
tcgen05.Field.SFB,
tCtSFB2_mma[sf_kblock_coord].iterator,
)
cute.gemm(
tiled_mma,
tCtAcc2,
tCrA[kblock_coord],
tCrB2[kblock_coord],
tCtAcc2,
)
# 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 < 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)
# Specialized epilogue warps
#
if warp_idx < self.mma_warp_id:
#
# Alloc tensor memory buffer
#
tmem.allocate(self.num_tmem_alloc_cols)
#
# Bar sync for retrieve tensor memory ptr from shared memory
#
tmem.wait_for_alloc()
#
# Retrieving tensor memory ptr and make accumulator tensor
#
acc1_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc1 = cute.make_tensor(acc1_tmem_ptr, tCtAcc_fake.layout)
tCtAcc2 = cute.make_tensor(acc1_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc1), tCtAcc_fake.layout)
#
# Partition for epilogue
#
epi_tidx = tidx
(
tiled_copy_t2r1,
tTR_tAcc1,
tTR_rAcc1,
) = self.epilog_tmem_copy_and_partition(
epi_tidx, tCtAcc1, tCgC, epi_tile, use_2cta_instrs
)
(
tiled_copy_t2r2,
tTR_tAcc2,
tTR_rAcc2,
) = self.epilog_tmem_copy_and_partition(
epi_tidx, tCtAcc2, tCgC, epi_tile, use_2cta_instrs
)
tTR_rC = cute.make_rmem_tensor(tTR_rAcc1.shape, self.c_dtype)
tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(
tiled_copy_t2r1, 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
)
#
# Get tile coord from block index (non-persistent: one tile per CTA)
#
block_idx = cute.arch.block_idx()
mma_tile_coord_mnl = (
block_idx[0] // cute.size(tiled_mma.thr_id.shape),
block_idx[1],
block_idx[2],
)
acc_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
)
#
# 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,
)
]
#
# Wait for accumulator buffer full
#
acc_pipeline.consumer_wait(acc_consumer_state)
tTR_tAcc1 = cute.group_modes(tTR_tAcc1, 3, cute.rank(tTR_tAcc1))
tTR_tAcc2 = cute.group_modes(tTR_tAcc2, 3, cute.rank(tTR_tAcc2))
bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
#
# Store accumulator to global memory in subtiles
#
subtile_cnt = cute.size(tTR_tAcc1.shape, mode=[3])
for subtile_idx in cutlass.range(subtile_cnt, unroll=1):
#
# Load accumulator from tensor memory buffer to register
#
tTR_tAcc1_mn = tTR_tAcc1[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r1, tTR_tAcc1_mn, tTR_rAcc1)
tTR_tAcc2_mn = tTR_tAcc2[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r2, tTR_tAcc2_mn, tTR_rAcc2)
acc_vec1 = tiled_copy_r2s.retile(tTR_rAcc1).load()
acc_vec2 = tiled_copy_r2s.retile(tTR_rAcc2).load()
acc_vec = cute.make_fragment(acc_vec1.shape, acc_vec1.dtype)
for i in cutlass.range_constexpr(0, cute.size(acc_vec.shape), 2):
# Compute silu(x) * y using optimized function
acc_vec[i], acc_vec[i + 1] = self.compute_silu_mul(
acc_vec1[i], acc_vec1[i + 1],
acc_vec2[i], acc_vec2[i + 1]
)
tRS_rC.store(acc_vec.load().to(self.c_dtype))
#
# Store C to shared memory
#
cute.copy(
tiled_copy_r2s,
tRS_rC,
tRS_sC[(None, None, None, subtile_idx)],
)
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, subtile_idx)],
bSG_gC[(None, subtile_idx)],
)
#
# Dealloc the tensor memory buffer
#
tmem.free(acc1_tmem_ptr)
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)
tAcc_epi = cute.flat_divide(
tAcc[((None, None), 0, 0)],
epi_tile,
)
# (EPI_TILE_M, EPI_TILE_N)
tiled_copy_t2r = tcgen05.make_tmem_copy(
copy_atom_t2r, tAcc_epi[(None, None, 0, 0)]
)
thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
# (T2R, T2R_M, T2R_N, EPI_M, EPI_M)
tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, 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, 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(
c: cute.Tensor,
cta_tile_shape_mnk: Tuple[int, int, int],
) -> Tuple[int, int, int]:
"""Compute simple grid: grid_m = M/tile_m, grid_n = N/tile_n.
:param c: The output tensor C
:type c: cute.Tensor
:param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile.
:type cta_tile_shape_mnk: tuple[int, int, int]
:return: Grid shape (num_tiles_m, num_tiles_n, batch_size) for kernel launch.
:rtype: tuple[int, int, int]
"""
c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0))
gc = cute.zipped_divide(c, tiler=c_shape)
num_ctas_mnl = gc[(0, (None, None, None))].shape
# grid = (num_tiles_m, num_tiles_n, batch_size)
grid = (num_ctas_mnl[0], num_ctas_mnl[1], num_ctas_mnl[2])
return grid
#### Do not change the code below this line ####
mma_tiler_mnk= (128, 128, 256)
# Stage numbers of shared memory and tmem
num_acc_stage = 1
num_ab_stage = 1
# Helper function for ceiling division
def ceil_div(a, b):
return (a + b - 1) // b
# GPU device kernel
@cute.kernel
def kernel(
tiled_mma: cute.TiledMma,
tma_atom_a: cute.CopyAtom,
mA_mkl: cute.Tensor,
tma_atom_b1: cute.CopyAtom,
mB_nkl1: cute.Tensor,
tma_atom_b2: cute.CopyAtom,
mB_nkl2: cute.Tensor,
tma_atom_sfa: cute.CopyAtom,
mSFA_mkl: cute.Tensor,
tma_atom_sfb1: cute.CopyAtom,
mSFB_nkl1: cute.Tensor,
tma_atom_sfb2: cute.CopyAtom,
mSFB_nkl2: cute.Tensor,
mC_mnl: cute.Tensor,
a_smem_layout_staged: cute.ComposedLayout,
b_smem_layout_staged: cute.ComposedLayout,
sfa_smem_layout_staged: cute.Layout,
sfb_smem_layout_staged: cute.Layout,
num_tma_load_bytes: cutlass.Constexpr[int],
epilogue_op: cutlass.Constexpr = lambda x: x
* (1.0 / (1.0 + cute.math.exp(-x, fastmath=True))),
):
"""
GPU device kernel performing the batched GEMM computation.
"""
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
tidx = cute.arch.thread_idx()
#
# 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)
# Coords outside cluster
cta_coord = (bidx, bidy, bidz)
mma_tile_coord_mnl = (
cta_coord[0] // cute.size(tiled_mma.thr_id.shape),
cta_coord[1],
cta_coord[2],
)
# Coord inside cta
tidx, _, _ = cute.arch.thread_idx()
#
# Define shared storage for kernel
#
@cute.struct
class SharedStorage:
ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_ab_stage * 2]
acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, num_acc_stage * 2]
tmem_holding_buf: cutlass.Int32
smem = utils.SmemAllocator()
storage = smem.allocate(SharedStorage)
# (MMA, MMA_M, MMA_K, STAGE)
sA = smem.allocate_tensor(
element_type=ab_dtype,
layout=a_smem_layout_staged.outer,
byte_alignment=128,
swizzle=a_smem_layout_staged.inner,
)
# (MMA, MMA_N, MMA_K, STAGE)
sB1 = smem.allocate_tensor(
element_type=ab_dtype,
layout=b_smem_layout_staged.outer,
byte_alignment=128,
swizzle=b_smem_layout_staged.inner,
)
# (MMA, MMA_N, MMA_K, STAGE)
sB2 = smem.allocate_tensor(
element_type=ab_dtype,
layout=b_smem_layout_staged.outer,
byte_alignment=128,
swizzle=b_smem_layout_staged.inner,
)
# (MMA, MMA_M, MMA_K, STAGE)
sSFA = smem.allocate_tensor(
element_type=sf_dtype,
layout=sfa_smem_layout_staged,
byte_alignment=128,
)
# (MMA, MMA_N, MMA_K, STAGE)
sSFB1 = smem.allocate_tensor(
element_type=sf_dtype,
layout=sfb_smem_layout_staged,
byte_alignment=128,
)
# (MMA, MMA_N, MMA_K, STAGE)
sSFB2 = smem.allocate_tensor(
element_type=sf_dtype,
layout=sfb_smem_layout_staged,
byte_alignment=128,
)
#
# Initialize mainloop ab_pipeline, acc_pipeline and their states
#
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1)
ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create(
barrier_storage=storage.ab_mbar_ptr.data_ptr(),
num_stages=num_ab_stage,
producer_group=ab_pipeline_producer_group,
consumer_group=ab_pipeline_consumer_group,
tx_count=num_tma_load_bytes,
).make_participants()
acc_producer, acc_consumer = pipeline.PipelineUmmaAsync.create(
barrier_storage=storage.acc_mbar_ptr.data_ptr(),
num_stages=num_acc_stage,
producer_group=ab_pipeline_producer_group,
consumer_group=pipeline.CooperativeGroup(
pipeline.Agent.Thread,
threads_per_cta,
),
).make_participants()
#
# Local_tile partition global tensors
#
# (bM, bK, RestM, RestK, RestL)
gA_mkl = cute.local_tile(
mA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
)
# (bN, bK, RestN, RestK, RestL)
gB_nkl1 = cute.local_tile(
mB_nkl1, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
)
# (bN, bK, RestN, RestK, RestL)
gB_nkl2 = cute.local_tile(
mB_nkl2, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
)
gSFA_mkl = cute.local_tile(
mSFA_mkl, cute.slice_(mma_tiler_mnk, (None, 0, None)), (None, None, None)
)
gSFB_nkl1 = cute.local_tile(
mSFB_nkl1, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
)
# (bN, bK, RestN, RestK, RestL)
gSFB_nkl2 = cute.local_tile(
mSFB_nkl2, cute.slice_(mma_tiler_mnk, (0, None, None)), (None, None, None)
)
# (bM, bN, RestM, RestN, RestL)
gC_mnl = cute.local_tile(
mC_mnl, cute.slice_(mma_tiler_mnk, (None, None, 0)), (None, None, None)
)
k_tile_cnt = cute.size(gA_mkl, mode=[3])
#
# Partition global tensor for TiledMMA_A/B/SFA/SFB/C
#
# (MMA, MMA_M, MMA_K, RestK)
thr_mma = tiled_mma.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)
tCgB1 = thr_mma.partition_B(gB_nkl1)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgB2 = thr_mma.partition_B(gB_nkl2)
# (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
tCgSFA = thr_mma.partition_A(gSFA_mkl)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgSFB1 = thr_mma.partition_B(gSFB_nkl1)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgSFB2 = thr_mma.partition_B(gSFB_nkl2)
# (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
tCgC = thr_mma.partition_C(gC_mnl)
#
# Partition global/shared tensor for TMA load A/B/SFA/SFB
#
# TMA Partition_S/D for A
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestM, RestK, RestL)
tAsA, tAgA = cpasync.tma_partition(
tma_atom_a,
0,
cute.make_layout(1),
cute.group_modes(sA, 0, 3),
cute.group_modes(tCgA, 0, 3),
)
# TMA Partition_S/D for B1
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestN, RestK, RestL)
tBsB1, tBgB1 = cpasync.tma_partition(
tma_atom_b1,
0,
cute.make_layout(1),
cute.group_modes(sB1, 0, 3),
cute.group_modes(tCgB1, 0, 3),
)
# TMA Partition_S/D for B2
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestN, RestK, RestL)
tBsB2, tBgB2 = cpasync.tma_partition(
tma_atom_b2,
0,
cute.make_layout(1),
cute.group_modes(sB2, 0, 3),
cute.group_modes(tCgB2, 0, 3),
)
# TMA Partition_S/D for SFA
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestM, RestK, RestL)
tAsSFA, tAgSFA = cpasync.tma_partition(
tma_atom_sfa,
0,
cute.make_layout(1),
cute.group_modes(sSFA, 0, 3),
cute.group_modes(tCgSFA, 0, 3),
)
tAsSFA = cute.filter_zeros(tAsSFA)
tAgSFA = cute.filter_zeros(tAgSFA)
# TMA Partition_S/D for SFB1
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestN, RestK, RestL)
tBsSFB1, tBgSFB1 = cpasync.tma_partition(
tma_atom_sfb1,
0,
cute.make_layout(1),
cute.group_modes(sSFB1, 0, 3),
cute.group_modes(tCgSFB1, 0, 3),
)
tBsSFB1 = cute.filter_zeros(tBsSFB1)
tBgSFB1 = cute.filter_zeros(tBgSFB1)
# TMA Partition_S/D for SFB2
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestN, RestK, RestL)
tBsSFB2, tBgSFB2 = cpasync.tma_partition(
tma_atom_sfb2,
0,
cute.make_layout(1),
cute.group_modes(sSFB2, 0, 3),
cute.group_modes(tCgSFB2, 0, 3),
)
tBsSFB2 = cute.filter_zeros(tBsSFB2)
tBgSFB2 = cute.filter_zeros(tBgSFB2)
#
# 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)
tCrB1 = tiled_mma.make_fragment_B(sB1)
# (MMA, MMA_N, MMA_K, STAGE)
tCrB2 = tiled_mma.make_fragment_B(sB2)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(mma_tiler_mnk[:2])
# (MMA, MMA_M, MMA_N)
tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
#
# Alloc tensor memory buffer
# Make ACC1 and ACC2 tmem tensor
# ACC1 += A @ B1
# ACC2 += A @ B2
#
tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=threads_per_cta,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
barrier_for_retrieve=tmem_alloc_barrier,
)
tmem.allocate(num_tmem_alloc_cols)
tmem.wait_for_alloc()
acc_tmem_ptr = tmem.retrieve_ptr(cutlass.Float32)
tCtAcc1 = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
acc_tmem_ptr1 = cute.recast_ptr(
acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc1),
dtype=cutlass.Float32,
)
tCtAcc2 = cute.make_tensor(acc_tmem_ptr1, tCtAcc_fake.layout)
#
# Make SFA/SFB1/SFB2 tmem tensor
#
# SFA tmem layout: (MMA, MMA_M, MMA_K)
tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
tiled_mma,
mma_tiler_mnk,
sf_vec_size,
cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
)
# Get SFA tmem ptr
sfa_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc1)
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc2),
dtype=sf_dtype,
)
tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
# SFB1, SFB2 tmem layout: (MMA, MMA_N, MMA_K)
tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
tiled_mma,
mma_tiler_mnk,
sf_vec_size,
cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
)
# Get SFB1 tmem ptr
sfb_tmem_ptr1 = cute.recast_ptr(
acc_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc1)
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc2)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA),
dtype=sf_dtype,
)
tCtSFB1 = cute.make_tensor(sfb_tmem_ptr1, tCtSFB_layout)
# Get SFB2 tmem ptr
sfb_tmem_ptr2 = cute.recast_ptr(
acc_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc1)
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc2)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFB1),
dtype=sf_dtype,
)
tCtSFB2 = cute.make_tensor(sfb_tmem_ptr2, tCtSFB_layout)
#
# Partition for S2T copy of SFA/SFB1/SFB2
#
# Make S2T CopyAtom
copy_atom_s2t = cute.make_copy_atom(
tcgen05.Cp4x32x128bOp(tcgen05.CtaGroup.ONE),
sf_dtype,
)
# (MMA, MMA_MN, MMA_K, STAGE)
tCsSFA_compact = cute.filter_zeros(sSFA)
# (MMA, MMA_MN, MMA_K)
tCtSFA_compact = cute.filter_zeros(tCtSFA)
tiled_copy_s2t_sfa = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFA_compact)
thr_copy_s2t_sfa = tiled_copy_s2t_sfa.get_slice(0)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFA_compact_s2t_ = thr_copy_s2t_sfa.partition_S(tCsSFA_compact)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFA_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
tiled_copy_s2t_sfa, tCsSFA_compact_s2t_
)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
tCtSFA_compact_s2t = thr_copy_s2t_sfa.partition_D(tCtSFA_compact)
# (MMA, MMA_MN, MMA_K, STAGE)
tCsSFB1_compact = cute.filter_zeros(sSFB1)
# (MMA, MMA_MN, MMA_K)
tCtSFB1_compact = cute.filter_zeros(tCtSFB1)
tiled_copy_s2t_sfb = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSFB1_compact)
thr_copy_s2t_sfb = tiled_copy_s2t_sfb.get_slice(0)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFB1_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB1_compact)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFB1_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
tiled_copy_s2t_sfb, tCsSFB1_compact_s2t_
)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
tCtSFB1_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB1_compact)
# SFB2 S2T copy and partition
# (MMA, MMA_MN, MMA_K, STAGE)
tCsSFB2_compact = cute.filter_zeros(sSFB2)
# (MMA, MMA_MN, MMA_K)
tCtSFB2_compact = cute.filter_zeros(tCtSFB2)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFB2_compact_s2t_ = thr_copy_s2t_sfb.partition_S(tCsSFB2_compact)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSFB2_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
tiled_copy_s2t_sfb, tCsSFB2_compact_s2t_
)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
tCtSFB2_compact_s2t = thr_copy_s2t_sfb.partition_D(tCtSFB2_compact)
#
# Slice to per mma tile index
#
# ((atom_v, rest_v), RestK)
tAgA = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
# ((atom_v, rest_v), RestK)
tBgB1 = tBgB1[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
# ((atom_v, rest_v), RestK)
tBgB2 = tBgB2[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
# ((atom_v, rest_v), RestK)
tAgSFA = tAgSFA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
# ((atom_v, rest_v), RestK)
tBgSFB1 = tBgSFB1[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
# ((atom_v, rest_v), RestK)
tBgSFB2 = tBgSFB2[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
#
# Execute Data copy and Math computation in the k_tile loop
#
if warp_idx == 0:
# Wait for accumulator buffer empty
acc_empty = acc_producer.acquire_and_advance()
# Set ACCUMULATE field to False for the first k_tile iteration
tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
# Execute k_tile loop
for k_tile in range(k_tile_cnt):
# Wait for AB buffer empty
ab_empty = ab_producer.acquire_and_advance()
# TMA load A/B1/B2/SFA/SFB1/SFB2 to shared memory
cute.copy(
tma_atom_a,
tAgA[(None, ab_empty.count)],
tAsA[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
)
cute.copy(
tma_atom_b1,
tBgB1[(None, ab_empty.count)],
tBsB1[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
)
cute.copy(
tma_atom_b2,
tBgB2[(None, ab_empty.count)],
tBsB2[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
)
cute.copy(
tma_atom_sfa,
tAgSFA[(None, ab_empty.count)],
tAsSFA[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
)
cute.copy(
tma_atom_sfb1,
tBgSFB1[(None, ab_empty.count)],
tBsSFB1[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
)
cute.copy(
tma_atom_sfb2,
tBgSFB2[(None, ab_empty.count)],
tBsSFB2[(None, ab_empty.index)],
tma_bar_ptr=ab_empty.barrier,
)
# Wait for AB buffer full
ab_full = ab_consumer.wait_and_advance()
# Copy SFA/SFB1/SFB2 to tmem
s2t_stage_coord = (None, None, None, None, ab_full.index)
tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
tCsSFB1_compact_s2t_staged = tCsSFB1_compact_s2t[s2t_stage_coord]
tCsSFB2_compact_s2t_staged = tCsSFB2_compact_s2t[s2t_stage_coord]
cute.copy(
tiled_copy_s2t_sfa,
tCsSFA_compact_s2t_staged,
tCtSFA_compact_s2t,
)
cute.copy(
tiled_copy_s2t_sfb,
tCsSFB1_compact_s2t_staged,
tCtSFB1_compact_s2t,
)
cute.copy(
tiled_copy_s2t_sfb,
tCsSFB2_compact_s2t_staged,
tCtSFB2_compact_s2t,
)
# tCtAcc1 += tCrA * tCrSFA * tCrB1 * tCrSFB1
# tCtAcc2 += tCrA * tCrSFA * tCrB2 * tCrSFB2
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_full.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,
tCtSFB1[sf_kblock_coord].iterator,
)
cute.gemm(
tiled_mma,
tCtAcc1,
tCrA[kblock_coord],
tCrB1[kblock_coord],
tCtAcc1,
)
tiled_mma.set(
tcgen05.Field.SFB,
tCtSFB2[sf_kblock_coord].iterator,
)
cute.gemm(
tiled_mma,
tCtAcc2,
tCrA[kblock_coord],
tCrB2[kblock_coord],
tCtAcc2,
)
# Enable accumulate on tCtAcc1/tCtAcc2 after first kblock
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
# Async arrive AB buffer empty
ab_full.release()
acc_empty.commit()
#
# Epilogue
# Partition for epilogue
#
op = tcgen05.Ld32x32bOp(tcgen05.Repetition.x128, tcgen05.Pack.NONE)
copy_atom_t2r = cute.make_copy_atom(op, cutlass.Float32)
tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tCtAcc1)
thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
# (T2R_M, T2R_N, EPI_M, EPI_M)
tTR_tAcc1 = thr_copy_t2r.partition_S(tCtAcc1)
# (T2R_M, T2R_N, EPI_M, EPI_M)
tTR_tAcc2 = thr_copy_t2r.partition_S(tCtAcc2)
# (T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL)
tTR_gC = thr_copy_t2r.partition_D(tCgC)
# (T2R_M, T2R_N, EPI_M, EPI_N）
tTR_rAcc1 = cute.make_rmem_tensor(
tTR_gC[None, None, None, None, 0, 0, 0].shape, cutlass.Float32
)
# (T2R_M, T2R_N, EPI_M, EPI_N）
tTR_rAcc2 = cute.make_rmem_tensor(
tTR_gC[None, None, None, None, 0, 0, 0].shape, cutlass.Float32
)
# (T2R_M, T2R_N, EPI_M, EPI_N）
tTR_rC = cute.make_rmem_tensor(
tTR_gC[None, None, None, None, 0, 0, 0].shape, c_dtype
)
# STG Atom
simt_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), c_dtype)
tTR_gC = tTR_gC[(None, None, None, None, *mma_tile_coord_mnl)]
# Wait for accumulator buffer full
acc_full = acc_consumer.wait_and_advance()
# Copy accumulator to register
cute.copy(tiled_copy_t2r, tTR_tAcc1, tTR_rAcc1)
cute.copy(tiled_copy_t2r, tTR_tAcc2, tTR_rAcc2)
# Silu activation on acc1 and multiply with acc2
acc_vec1 = epilogue_op(tTR_rAcc1.load())
acc_vec2 = tTR_rAcc2.load()
acc_vec = acc_vec1 * acc_vec2
tTR_rC.store(acc_vec.to(c_dtype))
# Store C to global memory
cute.copy(simt_atom, tTR_rC, tTR_gC)
acc_full.release()
# Deallocate TMEM
cute.arch.barrier()
tmem.free(acc_tmem_ptr)
return
@cute.jit
def my_kernel(
a_ptr: cute.Pointer,
b1_ptr: cute.Pointer,
b2_ptr: cute.Pointer,
sfa_ptr: cute.Pointer,
sfb1_ptr: cute.Pointer,
sfb2_ptr: cute.Pointer,
c_ptr: cute.Pointer,
problem_size: tuple,
epilogue_op: cutlass.Constexpr = lambda x: x
* (1.0 / (1.0 + cute.math.exp(-x, fastmath=True))),
):
"""
Host-side JIT function to prepare tensors and launch GPU kernel.
"""
m, n, k, l = problem_size
# Setup attributes that depend on gemm inputs
a_tensor = cute.make_tensor(
a_ptr,
cute.make_layout(
(m, cute.assume(k, 32), l),
stride=(cute.assume(k, 32), 1, cute.assume(m * k, 32)),
),
)
b_tensor1 = cute.make_tensor(
b1_ptr,
cute.make_layout(
(n, cute.assume(k, 32), l),
stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)),
),
)
b_tensor2 = cute.make_tensor(
b2_ptr,
cute.make_layout(
(n, cute.assume(k, 32), l),
stride=(cute.assume(k, 32), 1, cute.assume(n * k, 32)),
),
)
c_tensor = cute.make_tensor(
c_ptr, cute.make_layout((cute.assume(m, 32), n, l), stride=(n, 1, m * n))
)
# 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(
a_tensor.shape, sf_vec_size
)
sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
# ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
b_tensor1.shape, sf_vec_size
)
sfb_tensor1 = cute.make_tensor(sfb1_ptr, sfb_layout)
sfb_tensor2 = cute.make_tensor(sfb2_ptr, sfb_layout)
mma_op = tcgen05.MmaMXF4NVF4Op(
sf_dtype,
(mma_tiler_mnk[0], mma_tiler_mnk[1], mma_inst_shape_k),
tcgen05.CtaGroup.ONE,
tcgen05.OperandSource.SMEM,
)
tiled_mma = cute.make_tiled_mma(mma_op)
cluster_layout_vmnk = cute.tiled_divide(
cute.make_layout((1, 1, 1)),
(tiled_mma.thr_id.shape,),
)
# Compute A/B/SFA/SFB/C shared memory layout
a_smem_layout_staged = sm100_utils.make_smem_layout_a(
tiled_mma,
mma_tiler_mnk,
ab_dtype,
num_ab_stage,
)
# B1 and B2 have the same size thus share the same smem layout
b_smem_layout_staged = sm100_utils.make_smem_layout_b(
tiled_mma,
mma_tiler_mnk,
ab_dtype,
num_ab_stage,
)
sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
tiled_mma,
mma_tiler_mnk,
sf_vec_size,
num_ab_stage,
)
# SFB1 and SFB2 have the same size thus share the same smem layout
sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
tiled_mma,
mma_tiler_mnk,
sf_vec_size,
num_ab_stage,
)
atom_thr_size = cute.size(tiled_mma.thr_id.shape)
# Setup TMA for A
a_smem_layout = cute.slice_(a_smem_layout_staged, (None, None, None, 0))
tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
a_tensor,
a_smem_layout,
mma_tiler_mnk,
tiled_mma,
cluster_layout_vmnk .shape,
)
# Setup TMA for B1
b_smem_layout = cute.slice_(b_smem_layout_staged, (None, None, None, 0))
tma_atom_b1, tma_tensor_b1 = cute.nvgpu.make_tiled_tma_atom_B(
cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
b_tensor1,
b_smem_layout,
mma_tiler_mnk,
tiled_mma,
cluster_layout_vmnk .shape,
)
# Setup TMA for B2
tma_atom_b2, tma_tensor_b2 = cute.nvgpu.make_tiled_tma_atom_B(
cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
b_tensor2,
b_smem_layout,
mma_tiler_mnk,
tiled_mma,
cluster_layout_vmnk .shape,
)
# Setup TMA for SFA
sfa_smem_layout = cute.slice_(
sfa_smem_layout_staged , (None, None, None, 0)
)
tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
sfa_tensor,
sfa_smem_layout,
mma_tiler_mnk,
tiled_mma,
cluster_layout_vmnk .shape,
internal_type=cutlass.Int16,
)
# Setup TMA for SFB1
sfb_smem_layout = cute.slice_(
sfb_smem_layout_staged , (None, None, None, 0)
)
tma_atom_sfb1, tma_tensor_sfb1 = cute.nvgpu.make_tiled_tma_atom_B(
cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
sfb_tensor1,
sfb_smem_layout,
mma_tiler_mnk,
tiled_mma,
cluster_layout_vmnk .shape,
internal_type=cutlass.Int16,
)
# Setup TMA for SFB2
tma_atom_sfb2, tma_tensor_sfb2 = cute.nvgpu.make_tiled_tma_atom_B(
cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE),
sfb_tensor2,
sfb_smem_layout,
mma_tiler_mnk,
tiled_mma,
cluster_layout_vmnk .shape,
internal_type=cutlass.Int16,
)
# Compute TMA load bytes
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)
num_tma_load_bytes = (
a_copy_size + b_copy_size * 2 + sfa_copy_size + sfb_copy_size * 2
) * atom_thr_size
# Compute grid size
grid = (
cute.ceil_div(c_tensor.shape[0], mma_tiler_mnk[0]),
cute.ceil_div(c_tensor.shape[1], mma_tiler_mnk[1]),
c_tensor.shape[2],
)
# Launch the kernel.
kernel(
# MMA (Matrix Multiply-Accumulate) configuration
tiled_mma, # Tiled MMA object defining NVFP4 GEMM compute pattern
# TMA (Tensor Memory Accelerator) atoms and tensors for shared input matrix A
tma_atom_a, # TMA copy atom defining how to load A from global memory
tma_tensor_a, # Tensor descriptor for A matrix (m, k, l) - shared by both GEMMs
# TMA atoms and tensors for first B matrix (B1)
tma_atom_b1, # TMA copy atom defining how to load B1 from global memory
tma_tensor_b1, # Tensor descriptor for B1 matrix (n, k, l) - first GEMM
# TMA atoms and tensors for second B matrix (B2)
tma_atom_b2, # TMA copy atom defining how to load B2 from global memory
tma_tensor_b2, # Tensor descriptor for B2 matrix (n, k, l) - second GEMM
# TMA atoms and tensors for scale factor A (shared)
tma_atom_sfa, # TMA copy atom for loading scale factors for A
tma_tensor_sfa, # Tensor descriptor for SFA (block scale factors for A) - shared
# TMA atoms and tensors for scale factor B1
tma_atom_sfb1, # TMA copy atom for loading scale factors for B1
tma_tensor_sfb1, # Tensor descriptor for SFB1 (block scale factors for B1)
# TMA atoms and tensors for scale factor B2
tma_atom_sfb2, # TMA copy atom for loading scale factors for B2
tma_tensor_sfb2, # Tensor descriptor for SFB2 (block scale factors for B2)
# Output tensor C (stores both C1 and C2 results)
c_tensor, # Output tensor where both GEMM results will be stored (m, n, l)
# Shared memory layouts with staging for pipelined execution
a_smem_layout_staged, # Staged shared memory layout for A (includes stage dimension)
b_smem_layout_staged, # Staged shared memory layout for B1/B2 (includes stage dimension)
sfa_smem_layout_staged, # Staged shared memory layout for SFA (includes stage dimension)
sfb_smem_layout_staged, # Staged shared memory layout for SFB1/SFB2 (includes stage dimension)
# Pipeline synchronization parameter
num_tma_load_bytes, # Total bytes to load per TMA transaction (for barrier setup)
# Epilogue operation
epilogue_op, # Epilogue operation to apply to output (e.g., element-wise ops)
).launch(
grid=grid,
block=[threads_per_cta, 1, 1],
cluster=(1, 1, 1),
)
return
##### Do not change the code above this line ####
MMA_TILER_MN_MAP = {
(256, 4096, 7168, 1): (256, 64),
(512, 4096, 7168, 1): (256, 128),
(256, 3072, 4096, 1): (256, 64),
(512, 3072, 7168, 1): (256, 128),
}
CLUSTER_SHAPE_MN_MAP = {
(256, 4096, 7168, 1): (2, 2),
(512, 4096, 7168, 1): (2, 2),
(256, 3072, 4096, 1): (2, 1),
(512, 3072, 7168, 1): (4, 1),
}
# Global cache for compiled kernels (per problem shape)
_compiled_kernel_cache = {}
def compile_kernel(problem_shape: Tuple[int, int, int, int]):
"""
Compile the kernel once per problem shape and cache it.
This should be called before any timing measurements.
Returns:
The compiled kernel function
"""
global _compiled_kernel_cache
if problem_shape in _compiled_kernel_cache:
return _compiled_kernel_cache[problem_shape]
# Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
a_ptr = make_ptr(
ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16
)
b1_ptr = make_ptr(
ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16
)
b2_ptr = make_ptr(
ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16
)
c_ptr = make_ptr(
c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16
)
sfa_ptr = make_ptr(
sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32
)
sfb1_ptr = make_ptr(
sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32
)
sfb2_ptr = make_ptr(
sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32
)
if problem_shape in MMA_TILER_MN_MAP:
# Use optimized DualGemmKernel for benchmark shapes
mma_tiler_mn = MMA_TILER_MN_MAP[problem_shape]
cluster_shape_mn = CLUSTER_SHAPE_MN_MAP[problem_shape]
dual_gemm = DualGemmKernel(
sf_vec_size,
mma_tiler_mn,
cluster_shape_mn,
)
_compiled_kernel_cache[problem_shape] = cute.compile(
dual_gemm, a_ptr, b1_ptr, b2_ptr, sfa_ptr, sfb1_ptr, sfb2_ptr, c_ptr, *problem_shape
)
else:
# Use my_kernel for non-benchmark shapes
_compiled_kernel_cache[problem_shape] = cute.compile(
my_kernel, a_ptr, b1_ptr, b2_ptr, sfa_ptr, sfb1_ptr, sfb2_ptr, c_ptr, problem_shape
)
return _compiled_kernel_cache[problem_shape]
def custom_kernel(data: input_t) -> output_t:
"""
Execute the block-scaled dual GEMM kernel with silu activation,
C = silu(A @ B1) * (A @ B2).
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 (a, b1, b2, sfa_cpu, sfb1_cpu, sfb2_cpu, c) PyTorch tensors
a: [m, k, l] - Input matrix in float4e2m1fn
b1: [n, k, l] - Input matrix in float4e2m1fn
b2: [n, k, l] - Input matrix in float4e2m1fn
sfa_cpu: [m, k, l] - Scale factors in float8_e4m3fn, used by reference implementation
sfb1_cpu: [n, k, l] - Scale factors in float8_e4m3fn, used by reference implementation
sfb2_cpu: [n, k, l] - Scale factors in float8_e4m3fn, used by reference implementation
sfa_permuted: [32, 4, rest_m, 4, rest_k, l] - Scale factors in float8_e4m3fn
sfb1_permuted: [32, 4, rest_n, 4, rest_k, l] - Scale factors in float8_e4m3fn
sfb2_permuted: [32, 4, rest_n, 4, rest_k, l] - Scale factors in float8_e4m3fn
c: [m, n, l] - Output vector in float16
Returns:
Output tensor c with computed results
"""
a, b1, b2, _, _, _, sfa_permuted, sfb1_permuted, sfb2_permuted, c = data
# Get dimensions from MxKxL layout
_, k, _ = a.shape
m, n, l = c.shape
# Torch use e2m1_x2 data type, thus k is halved
k = k * 2
problem_shape = (m, n, k, l)
# Ensure kernel is compiled (will use cached version if available)
# To avoid the compilation overhead, we compile the kernel once and cache it.
compiled_func = compile_kernel(problem_shape)
# Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
a_ptr = make_ptr(
ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
)
b1_ptr = make_ptr(
ab_dtype, b1.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
)
b2_ptr = make_ptr(
ab_dtype, b2.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
)
c_ptr = make_ptr(
c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
)
sfa_ptr = make_ptr(
sf_dtype, sfa_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
sfb1_ptr = make_ptr(
sf_dtype, sfb1_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
sfb2_ptr = make_ptr(
sf_dtype, sfb2_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
# Execute the compiled kernel
# DualGemmKernel has problem size baked in at compile time (7 args)
# my_kernel needs problem_size passed at runtime (8 args)
if problem_shape in MMA_TILER_MN_MAP:
compiled_func(a_ptr, b1_ptr, b2_ptr, sfa_ptr, sfb1_ptr, sfb2_ptr, c_ptr)
else:
compiled_func(a_ptr, b1_ptr, b2_ptr, sfa_ptr, sfb1_ptr, sfb2_ptr, c_ptr, problem_shape)
return cscrolls · 2659 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 382879.
⋯ diff truncated: revisions differ almost entirely
Best evidence level for this revision: reported
JSON