submission 382560
Arseni Ivanov · python · License unknown
Use it
Vendorable · source mirrored · license unknownView source →
No package. Vendor the mirrored source: 2613 lines, June 9 Researcher Reciprocity License v1.0.
temp.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-modal-nvfp4-dual-gemm-382560?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:ab9dad907a5b153b275538bbde3eb293cd4279b1befb9f900e13bf330347cbba
license declaredunknown
license concludedunknown
authorsArseni Ivanov
imported2026-08-15
Techniques
Extracted from the mirrored source by pattern, never inferred. Each row cites its line.
fp4
) #fp8 = (128, 64, 32), fp4 = (128, 64, 64)fused-epilogue
self.epi_tile = sm100_utils.compute_epilogue_tile_shape(mbarrier
self.epilog_sync_barrier = pipeline.NamedBarrier(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
temp.py2613 lines
#!POPCORN leaderboard modal_nvfp4_dual_gemm
#!POPCORN gpu B200
import cutlass
import cutlass.cute as cute
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.pipeline as pipeline
import cutlass.utils as utils
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.cute.runtime import make_ptr
from typing import Tuple, Type, Union
from task import input_t, output_t
from cutlass.cute.tensor import TensorSSA
from cutlass import Float32, Float16, Int8, Int32, Int64
from cutlass.cutlass_dsl import T, dsl_user_op
from cutlass._mlir import ir
from cutlass._mlir.dialects import nvvm, arith, llvm, vector, builtin
@dsl_user_op
def swiglu_f32(
acc_gate_vec: ir.Value,
acc_val_vec: ir.Value,
*,
loc=None,
ip=None
) -> ir.Value:
"""
Computes SwiGLU: Res = (Gate * Sigmoid(Gate)) * Value in FP32, then packs to FP16.
Fixed Register Indices:
$0 = Output (Packed f16x2)
$1 = Gate[0]
$2 = Gate[1]
$3 = Value[0]
$4 = Value[1]
"""
vec_type = ir.VectorType(acc_gate_vec.type)
num_elements = vec_type.shape[0]
num_pairs = num_elements // 2
# Output: vector<N x f16>
out_f16_type = ir.VectorType.get([num_elements], Float16.mlir_type, loc=loc)
# Intermediate: vector<N/2 x i32>
packed_vec_type = ir.VectorType.get([num_pairs], Int32.mlir_type, loc=loc)
packed_res = llvm.mlir_undef(packed_vec_type, loc=loc, ip=ip)
asm_str = """
{
.reg .f32 %half_g0, %half_g1;
.reg .f32 %tanh_in1;
.reg .f32 %tanh0, %tanh1;
mul.f32 %half_g0, $1, 0.5;
tanh.approx.f32 %tanh0, %half_g0;
// additive bias to fix single precision error
mul.f32 %half_g1, $2, 0.5;
fma.rn.f32 %tanh_in1, $2, 0.5, 0.0002; //redo gate computation to avoid i/o depedendency
tanh.approx.f32 %tanh1, %tanh_in1;
fma.rn.f32 %tanh0, %half_g0, %tanh0, %half_g0;
fma.rn.f32 %tanh_in1, %half_g1, %tanh1, %half_g1;
mul.f32 %tanh0, %tanh0, $3;
mul.f32 %tanh_in1, %tanh_in1, $4;
cvt.rn.f16x2.f32 $0, %tanh_in1, %tanh0;
}
"""
for i in range(num_pairs):
idx0 = arith.constant(Int32.mlir_type, i * 2, loc=loc, ip=ip)
idx1 = arith.constant(Int32.mlir_type, i * 2 + 1, loc=loc, ip=ip)
g0 = llvm.extractelement(acc_gate_vec, idx0, loc=loc, ip=ip)
g1 = llvm.extractelement(acc_gate_vec, idx1, loc=loc, ip=ip)
v0 = llvm.extractelement(acc_val_vec, idx0, loc=loc, ip=ip)
v1 = llvm.extractelement(acc_val_vec, idx1, loc=loc, ip=ip)
packed_val = llvm.inline_asm(
Int32.mlir_type,
[g0, g1, v0, v1],
asm_str,
"=r,f,f,f,f",
has_side_effects=False,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
ins_idx = arith.constant(Int32.mlir_type, i, loc=loc, ip=ip)
packed_res = llvm.insertelement(packed_res, packed_val, ins_idx, loc=loc, ip=ip)
# 3. Bitcast to vector<N x f16>
res_f16 = llvm.bitcast(out_f16_type, packed_res, loc=loc, ip=ip)
return res_f16
#mma_tiler_mnk= (128, 64, 256)
#mma_inst_shape_k = 64
ab_dtype = cutlass.Float4E2M1FN
sf_dtype = cutlass.Float8E4M3FN
c_dtype = cutlass.Float16
sf_vec_size = 16
threads_per_cta = 128
OCCUPANCY = 1
_TMA_CACHE_EVICT_NORMAL = 0x1000000000000000
_TMA_CACHE_EVICT_FIRST = 0x12F0000000000000
_TMA_CACHE_EVICT_LAST = 0x14F0000000000000
class Sm100BlockScaledSmall:
def __init__(
self,
sf_vec_size: int,
mma_tiler_mn: Tuple[int, int],
cluster_shape_mn: Tuple[int, int],
occupancy: int = 1,
):
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.mma_tiler = (*mma_tiler_mn, 1)
self.cta_group = (
tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
)
self.occupancy = int(occupancy)
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)
)
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")
self.num_tmem_alloc_cols = 512
def _setup_attributes(self):
self.mma_inst_shape_mn = (
self.mma_tiler[0],
self.mma_tiler[1],
)
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,
)
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],
)
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,),
)
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
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 = 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,
)
self.prefetch_stage = self.num_ab_stage
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,
)
self.overlapping_accum = self.num_acc_stage == 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.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k
self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols
self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] * self.num_acc_stage if not self.overlapping_accum else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols
self.iter_acc_early_release_in_epilogue = self.num_sf_tmem_cols // self.epi_tile_n
@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,
problem_size: tuple,
max_active_clusters: cutlass.Constexpr,
epilogue_op: cutlass.Constexpr = lambda x: x
* (1.0 / (1.0 + cute.math.exp(-x, fastmath=True))),
):
m, n, k, l = problem_size
sf_k = k // self.sf_vec_size
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)),
),
)
b1_tensor = 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)),
),
)
b2_tensor = 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)),
)
sfa_tensor = cute.make_tensor(
sfa_ptr,
cute.make_layout((m, sf_k, l), stride=(sf_k, 1, m * sf_k)),
)
sfb1_tensor = cute.make_tensor(
sfb1_ptr,
cute.make_layout((n, sf_k, l), stride=(sf_k, 1, n * sf_k)),
)
sfb2_tensor = cute.make_tensor(
sfb2_ptr,
cute.make_layout((n, sf_k, l), stride=(sf_k, 1, n * sf_k)),
)
a_tensor.mark_compact_shape_dynamic(mode=1, stride_order=(2, 0, 1), divisibility=2)
b1_tensor.mark_compact_shape_dynamic(mode=1, stride_order=(2, 0, 1), divisibility=2)
b2_tensor.mark_compact_shape_dynamic(mode=1, stride_order=(2, 0, 1), divisibility=2)
c_tensor.mark_compact_shape_dynamic(mode=1, stride_order=(2, 0, 1), divisibility=1)
self.a_dtype: Type[cutlass.Numeric] = a_tensor.element_type
self.b_dtype: Type[cutlass.Numeric] = b1_tensor.element_type
self.sf_dtype: Type[cutlass.Numeric] = sfa_tensor.element_type
self.c_dtype: Type[cutlass.Numeric] = c_tensor.element_type
self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode()
self.b_major_mode = utils.LayoutEnum.from_tensor(b1_tensor).mma_major_mode()
self.c_layout = utils.LayoutEnum.from_tensor(c_tensor)
self._setup_attributes()
sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
a_tensor.shape, self.sf_vec_size
)
sfa_tensor = cute.make_tensor(sfa_tensor.iterator, sfa_layout)
sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(
b1_tensor.shape, self.sf_vec_size
)
sfb1_tensor = cute.make_tensor(sfb1_tensor.iterator, sfb_layout)
sfb2_tensor = cute.make_tensor(sfb2_tensor.iterator, 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)
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,
)
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,
)
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,
)
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
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,
)
grid = self._compute_grid(c_tensor, self.cta_tile_shape_mnk, self.cluster_shape_mn)
self.buffer_align_bytes = 1024
@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
sC: cute.struct.Align[
cute.struct.MemRange[
self.c_dtype,
cute.cosize(self.c_smem_layout_staged.outer),
],
self.buffer_align_bytes,
]
sA: cute.struct.Align[
cute.struct.MemRange[
self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)
],
self.buffer_align_bytes,
]
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,
]
sSFA: cute.struct.Align[
cute.struct.MemRange[
self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)
],
self.buffer_align_bytes,
]
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
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,
epilogue_op,
).launch(
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
min_blocks_per_mp=self.occupancy,
smem=self.shared_storage.size_in_bytes(),
)
return
@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_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: cute.CopyAtom,
mC_mnl: cute.Tensor,
cluster_layout_vmnk: cute.Layout,
cluster_layout_sfb_vmnk: cute.Layout,
a_smem_layout_staged: cute.ComposedLayout,
b_smem_layout_staged: cute.ComposedLayout,
sfa_smem_layout_staged: cute.Layout,
sfb_smem_layout_staged: cute.Layout,
c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout],
epi_tile: cute.Tile,
epilogue_op: cutlass.Constexpr,
):
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2
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
)
tidx, _, _ = cute.arch.thread_idx()
smem = utils.SmemAllocator()
storage = smem.allocate(self.shared_storage)
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,
)
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,
)
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,
)
cute.arch.cluster_arrive_relaxed()
sC = storage.sC.get_tensor(
c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner
)
sA = storage.sA.get_tensor(
a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner
)
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
)
sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
sSFB1 = storage.sSFB1.get_tensor(sfb_smem_layout_staged)
sSFB2 = storage.sSFB2.get_tensor(sfb_smem_layout_staged)
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
)
gA_mkl = cute.local_tile(
mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
gB1_nkl = cute.local_tile(
mB1_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
gB2_nkl = cute.local_tile(
mB2_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
gSFA_mkl = cute.local_tile(
mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
gSFB1_nkl = cute.local_tile(
mSFB1_nkl,
cute.slice_(self.mma_tiler_sfb, (0, None, None)),
(None, None, None),
)
gSFB2_nkl = cute.local_tile(
mSFB2_nkl,
cute.slice_(self.mma_tiler_sfb, (0, None, None)),
(None, None, None),
)
gC_mnl = cute.local_tile(
mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
)
k_block_cnt = cute.size(gA_mkl, mode=[3])
thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v)
tCgA = thr_mma.partition_A(gA_mkl)
tCgB1 = thr_mma.partition_B(gB1_nkl)
tCgB2 = thr_mma.partition_B(gB2_nkl)
tCgSFA = thr_mma.partition_A(gSFA_mkl)
tCgSFB1 = thr_mma_sfb.partition_B(gSFB1_nkl)
tCgSFB2 = thr_mma_sfb.partition_B(gSFB2_nkl)
tCgC = thr_mma.partition_C(gC_mnl)
a_cta_layout = cute.make_layout(
cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
)
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),
)
b_cta_layout = cute.make_layout(
cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
)
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),
)
sfa_cta_layout = a_cta_layout
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)
sfb_cta_layout = cute.make_layout(
cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape
)
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),
)
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),
)
tBgSFB1 = cute.filter_zeros(tBgSFB1)
tBgSFB2 = cute.filter_zeros(tBgSFB2)
tBsSFB1 = cute.filter_zeros(tBsSFB1)
tBsSFB2 = cute.filter_zeros(tBsSFB2)
tCrA = tiled_mma.make_fragment_A(sA)
tCrB1 = tiled_mma.make_fragment_B(sB1)
tCrB2 = tiled_mma.make_fragment_B(sB2)
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)
)
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:
tCtAcc_fake = tiled_mma.make_fragment_C(
cute.append(acc_shape, self.num_acc_stage)
)
cute.arch.cluster_wait()
if warp_idx == self.tma_warp_id:
ab_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_ab_stage
)
mma_tile_coord_mnl = (
bidx // cute.size(tiled_mma.thr_id.shape),
bidy,
bidz,
)
tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
tBgB1_slice = tBgB1[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
tBgB2_slice = tBgB2[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
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
tBgSFB1_slice = tBgSFB1[(None, slice_n, None, mma_tile_coord_mnl[2])]
tBgSFB2_slice = tBgSFB2[(None, slice_n, None, mma_tile_coord_mnl[2])]
ab_producer_state.reset_count()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < k_block_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state)
for k_block_idx in cutlass.range(0, k_block_cnt, 1, unroll=1):
ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status)
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,
)
ab_producer_state.advance()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < k_block_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state)
ab_pipeline.producer_tail(ab_producer_state)
if warp_idx == self.mma_warp_id:
tmem.wait_for_alloc()
acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
sfa_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr + self.num_accumulator_tmem_cols,
dtype=self.sf_dtype,
)
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)
#sus, might need to make a pointer with other b
sfb1_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,
)
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(
acc_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc_base)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFB1),
dtype=self.sf_dtype,
)
tCtSFB2 = cute.make_tensor(sfb2_tmem_ptr, tCtSFB_layout)
(
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)
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
)
mma_tile_coord_mnl = (
bidx // cute.size(tiled_mma.thr_id.shape),
bidy,
bidz,
)
if cutlass.const_expr(self.overlapping_accum):
acc_stage_index = acc_producer_state.phase ^ 1
else:
acc_stage_index = acc_producer_state.index
tCtAcc1 = tCtAcc_base[(None, None, None, acc_stage_index)]
tCtAcc2 = tCtAcc_base[(None, None, None, acc_stage_index + 1)]
ab_consumer_state.reset_count()
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < k_block_cnt and is_leader_cta:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
)
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):
pass
elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
# Adjust SFB1
shifted_ptr1 = cute.recast_ptr(
acc_tmem_ptr
+ self.num_accumulator_tmem_cols
+ self.num_sfa_tmem_cols
+ offset,
dtype=self.sf_dtype,
)
tCtSFB1_mma = cute.make_tensor(shifted_ptr1, tCtSFB_layout)
# Adjust SFB2
shifted_ptr2 = cute.recast_ptr(
acc_tmem_ptr
+ self.num_accumulator_tmem_cols
+ self.num_sfa_tmem_cols
+ self.num_sfb_tmem_cols # Add stride to get to SFB2
+ offset, # Add the alternating offset
dtype=self.sf_dtype,
)
tCtSFB2_mma = cute.make_tensor(shifted_ptr2, tCtSFB_layout)
tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
for k_tile in range(k_block_cnt):
if is_leader_cta:
ab_pipeline.consumer_wait(
ab_consumer_state, peek_ab_full_status
)
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,
)
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,
)
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,
)
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
ab_pipeline.consumer_release(ab_consumer_state)
ab_consumer_state.advance()
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < k_block_cnt:
if is_leader_cta:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
)
if is_leader_cta:
acc_pipeline.producer_commit(acc_producer_state)
acc_producer_state.advance()
acc_pipeline.producer_tail(acc_producer_state)
if warp_idx < self.mma_warp_id:
tmem.allocate(self.num_tmem_alloc_cols)
tmem.wait_for_alloc()
acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
epi_tidx = tidx
(
tiled_copy_t2r,
tTR_tAcc1_base,
tTR_rAcc1,
) = self.epilog_tmem_copy_and_partition(
epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs
)
(
tiled_copy_t2r,
tTR_tAcc2_base,
tTR_rAcc2,
) = self.epilog_tmem_copy_and_partition(
epi_tidx, tCtAcc_base, 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_t2r, tTR_rC, epi_tidx, sC
)
(
tma_atom_c,
bSG_sC,
bSG_gC_partitioned,
) = self.epilog_gmem_copy_and_partition(
epi_tidx, tma_atom_c, tCgC, epi_tile, sC
)
acc_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
)
mma_tile_coord_mnl = (
bidx // cute.size(tiled_mma.thr_id.shape),
bidy,
bidz,
)
bSG_gC = bSG_gC_partitioned[
(
None,
None,
None,
*mma_tile_coord_mnl,
)
]
acc_stage_index = acc_consumer_state.index
tTR_tAcc1 = tTR_tAcc1_base[
(None, None, None, None, None, acc_stage_index)
]
tTR_tAcc2 = tTR_tAcc2_base[
(None, None, None, None, None, acc_stage_index+1)
]
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))
subtile_cnt = cute.size(tTR_tAcc1.shape, mode=[3])
for subtile_idx in cutlass.range(subtile_cnt):
tTR_tAcc_mn_2 = tTR_tAcc2[(None, None, None, subtile_idx)]
tTR_tAcc_mn_1 = tTR_tAcc1[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r, tTR_tAcc_mn_1, tTR_rAcc1)
cute.copy(tiled_copy_t2r, tTR_tAcc_mn_2, tTR_rAcc2)
acc1_vec = tTR_rAcc1.load()
acc2_vec = tTR_rAcc2.load()
acc2_vec = swiglu_f32(acc1_vec, acc2_vec)
acc2_vec = TensorSSA(acc2_vec, tRS_rC.layout, cutlass.Float16)
tRS_rC.store(acc2_vec)
cute.copy(
tiled_copy_r2s,
tRS_rC,
tRS_sC[(None, None, None, subtile_idx)],
)
if warp_idx == self.epilog_warp_id[0]:
cute.copy(
tma_atom_c,
bSG_sC[(None, subtile_idx)],
bSG_gC[(None, subtile_idx)],
)
self.epilog_sync_barrier.arrive_and_wait()
with cute.arch.elect_one():
acc_pipeline.consumer_release(acc_consumer_state)
acc_consumer_state.advance()
tmem.relinquish_alloc_permit()
tmem.free(acc_tmem_ptr)
def mainloop_s2t_copy_and_partition(
self,
sSF: cute.Tensor,
tSF: cute.Tensor,
) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
tCsSF_compact = cute.filter_zeros(sSF)
tCtSF_compact = cute.filter_zeros(tSF)
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)
tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
tiled_copy_s2t, tCsSF_compact_s2t_
)
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]:
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,
)
tAcc_epi = cute.flat_divide(
tAcc[((None, None), 0, 0, None)],
epi_tile,
)
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)
tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
gC_mnl_epi = cute.flat_divide(
gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
)
tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
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]:
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)
thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
tRS_sC = thr_copy_r2s.partition_D(sC)
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]:
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)
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]:
num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2
num_c_stage = 1
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) * 2
+ cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one)
+ cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) * 2
)
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
num_ab_stage = (
smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes)
) // ab_bytes_per_stage
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],
cluster_shape_mn: Tuple[int, int],
) -> Tuple[int, int, int]:
"""Compute grid shape for the output tensor C.
: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]
:param cluster_shape_mn: Shape of each cluster in M, N dimensions.
:type cluster_shape_mn: tuple[int, int]
:return: Grid shape for kernel launch.
:rtype: tuple[int, int, int]
"""
grid = (
cute.ceil_div(c.layout.shape[0], cta_tile_shape_mnk[0]),
cute.ceil_div(c.layout.shape[1], cta_tile_shape_mnk[1]),
c.layout.shape[2],
)
return grid
class Sm100BlockScaledDenseGemmKernel:
def __init__(
self,
sf_vec_size: int, #always 16
mma_tiler_mn: Tuple[int, int], #(128, 64) or (128, 128)
cluster_shape_mn: Tuple[int, int], #(1, 2)
occupancy: int = 1, #always 1
):
self.acc_dtype = cutlass.Float32
self.sf_vec_size = sf_vec_size
self.cluster_shape_mn = cluster_shape_mn
self.mma_tiler = (*mma_tiler_mn, 1)
self.cta_group = (
tcgen05.CtaGroup.ONE
)
self.occupancy = int(occupancy)
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) #6 * 32 = 192
)
self.epilog_sync_barrier = pipeline.NamedBarrier( #32 * 4 = 128 epilog barrier threads
barrier_id=1,
num_threads=32 * len(self.epilog_warp_id),
)
self.tmem_alloc_barrier = pipeline.NamedBarrier( #32 * 5 = 160 tmem barrier threads
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") #227*1024 b = 232 kB
self.num_tmem_alloc_cols = 256
def _setup_attributes(self):
self.mma_inst_shape_mn = ( #(128,64) or (128, 128)
self.mma_tiler[0],
self.mma_tiler[1],
)
self.mma_inst_shape_mn_sfb = (
self.mma_inst_shape_mn[0],
cute.round_up(self.mma_inst_shape_mn[1], 128), #always (128, 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,
) #fp8 = (128, 64, 32), fp4 = (128, 64, 64)
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,
) #(128, 128, 32)?
mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
mma_inst_tile_k = 4
self.mma_tiler = ( #(128, 64, 256)
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,
) #(128, 128, 256)
self.cluster_layout_vmnk = cute.tiled_divide(
cute.make_layout((*self.cluster_shape_mn, 1)),
(tiled_mma.thr_id.shape,),
) #((1), (1,2,1): (0), (0,1,0))
self.cluster_layout_sfb_vmnk = cute.tiled_divide(
cute.make_layout((*self.cluster_shape_mn, 1)),
(tiled_mma_sfb.thr_id.shape,),
)
self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) #2
self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) #1
self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) #1
self.is_a_mcast = self.num_mcast_ctas_a > 1 #true
self.is_b_mcast = self.num_mcast_ctas_b > 1 #false
self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 #false
self.epi_tile = sm100_utils.compute_epilogue_tile_shape( #(128,32)
self.mma_tiler,
False,
self.c_layout,
self.c_dtype,
)
self.num_acc_stage, self.num_ab_stage, self.num_c_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,
) #2, 4/5, 1
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,
)
sf_atom_mn = 32
self.num_accumulator_tmem_cols = self.mma_tiler[1] * self.num_acc_stage #64 * 2 or 128 * 2
self.total_sfa_cols = self.num_accumulator_tmem_cols + (self.mma_tiler[0] // sf_atom_mn) * mma_inst_tile_k #128 + 4*4 = 144
self.total_sfb_cols = self.total_sfa_cols + (self.mma_tiler_sfb[1] // sf_atom_mn) * mma_inst_tile_k #144 + 2*4 or 4*4 = 160
@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,
problem_size: tuple,
):
m, n, k, l = problem_size
a_tensor = cute.make_tensor(
a_ptr,
cute.make_ordered_layout(
(cute.assume(m, 32), k, l), order=(1, 0, 2)
),
)
b1_tensor = cute.make_tensor(
b1_ptr,
cute.make_ordered_layout(
(cute.assume(n, 32), k, l), order=(1, 0, 2)
),
)
b2_tensor = cute.make_tensor(
b2_ptr,
cute.make_ordered_layout(
(cute.assume(n, 32), k, l), order=(1, 0, 2)
),
)
c_tensor = cute.make_tensor(
c_ptr,
cute.make_ordered_layout(
(m, cute.assume(n, 32), l), order=(1, 0, 2)
)
)
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)
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)
self.a_dtype: Type[cutlass.Numeric] = a_tensor.element_type
self.b_dtype: Type[cutlass.Numeric] = b1_tensor.element_type
self.sf_dtype: Type[cutlass.Numeric] = sfa_tensor.element_type
self.c_dtype: Type[cutlass.Numeric] = c_tensor.element_type
self.a_major_mode, self.b_major_mode, self.c_layout = (
tcgen05.OperandMajorMode.K,
tcgen05.OperandMajorMode.K,
utils.LayoutEnum.ROW_MAJOR,
)
self._setup_attributes()
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)
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,
)
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,
)
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,
)
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
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,
)
grid_m = (m // self.mma_tiler[0]) * cute.size(tiled_mma.thr_id.shape)
grid_n = n // self.mma_tiler[1]
grid = (grid_m, grid_n, l) #(m//128, n//128 or 64, 1)
self.buffer_align_bytes = 128
@cute.struct
class SharedStorage:
ab_full_mbar_ptr: cute.struct.Align[
cute.struct.MemRange[cutlass.Int64, self.num_ab_stage],
16
]
ab_empty_mbar_ptr: cute.struct.Align[
cute.struct.MemRange[cutlass.Int64, self.num_ab_stage],
16
]
acc_full_mbar_ptr: cute.struct.Align[
cute.struct.MemRange[cutlass.Int64, self.num_acc_stage],
16
]
acc_empty_mbar_ptr: cute.struct.Align[
cute.struct.MemRange[cutlass.Int64, self.num_acc_stage],
16
]
tmem_holding_buf: cutlass.Int32
sC: cute.struct.Align[ #8kb
cute.struct.MemRange[
self.c_dtype,
cute.cosize(self.c_smem_layout_staged.outer),
],
self.buffer_align_bytes,
]
sA: cute.struct.Align[ #16kb
cute.struct.MemRange[
self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)
],
self.buffer_align_bytes,
]
sB1: cute.struct.Align[ #8kb
cute.struct.MemRange[
self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)
],
self.buffer_align_bytes,
]
sB2: cute.struct.Align[ #8kb
cute.struct.MemRange[
self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)
],
self.buffer_align_bytes,
]
sSFA: cute.struct.Align[ #2kb
cute.struct.MemRange[
self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)
],
self.buffer_align_bytes,
]
sSFB1: cute.struct.Align[ #1kb
cute.struct.MemRange[
self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)
],
self.buffer_align_bytes,
]
sSFB2: cute.struct.Align[ #1kb
cute.struct.MemRange[
self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)
],
self.buffer_align_bytes,
]
self.shared_storage = SharedStorage
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),
min_blocks_per_mp=self.occupancy,
smem=self.shared_storage.size_in_bytes()
)
return
@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_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: cute.CopyAtom,
mC_mnl: cute.Tensor,
cluster_layout_vmnk: cute.Layout,
cluster_layout_sfb_vmnk: cute.Layout,
a_smem_layout_staged: cute.ComposedLayout,
b_smem_layout_staged: cute.ComposedLayout,
sfa_smem_layout_staged: cute.Layout,
sfb_smem_layout_staged: cute.Layout,
c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout],
epi_tile: cute.Tile,
):
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
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)
bidx, bidy, bidz = cute.arch.block_idx()
mma_tile_coord_mnl = (
bidx,
bidy,
bidz,
)
tidx, _, _ = cute.arch.thread_idx()
smem = utils.SmemAllocator()
storage = smem.allocate(self.shared_storage)
mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
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,
)
acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
num_acc_consumer_threads = len(self.epilog_warp_id)
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,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
barrier_for_retrieve=self.tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
)
sC = storage.sC.get_tensor(
c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner
)
sA = storage.sA.get_tensor(
a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner
)
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
)
sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
sSFB1 = storage.sSFB1.get_tensor(sfb_smem_layout_staged)
sSFB2 = storage.sSFB2.get_tensor(sfb_smem_layout_staged)
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
)
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):
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
)
gA_mkl = cute.local_tile(
mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
gB1_nkl = cute.local_tile(
mB1_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
gB2_nkl = cute.local_tile(
mB2_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
gSFA_mkl = cute.local_tile(
mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
gSFB1_nkl = cute.local_tile(
mSFB1_nkl,
cute.slice_(self.mma_tiler_sfb, (0, None, None)),
(None, None, None),
)
gSFB2_nkl = cute.local_tile(
mSFB2_nkl,
cute.slice_(self.mma_tiler_sfb, (0, None, None)),
(None, None, None),
)
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])
thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v)
tCgA = thr_mma.partition_A(gA_mkl)
tCgB1 = thr_mma.partition_B(gB1_nkl)
tCgB2 = thr_mma.partition_B(gB2_nkl)
tCgSFA = thr_mma.partition_A(gSFA_mkl)
tCgSFB1 = thr_mma_sfb.partition_B(gSFB1_nkl)
tCgSFB2 = thr_mma_sfb.partition_B(gSFB2_nkl)
tCgC = thr_mma.partition_C(gC_mnl)
a_cta_layout = cute.make_layout(
cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
)
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),
)
b_cta_layout = cute.make_layout(
cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
)
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),
)
sfa_cta_layout = a_cta_layout
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)
sfb_cta_layout = cute.make_layout(
cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape
)
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),
)
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),
)
tBgSFB1 = cute.filter_zeros(tBgSFB1)
tBgSFB2 = cute.filter_zeros(tBgSFB2)
tBsSFB1 = cute.filter_zeros(tBsSFB1)
tBsSFB2 = cute.filter_zeros(tBsSFB2)
tCrA = tiled_mma.make_fragment_A(sA)
tCrB1 = tiled_mma.make_fragment_B(sB1)
tCrB2 = tiled_mma.make_fragment_B(sB2)
acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
tCtAcc_fake = tiled_mma.make_fragment_C(
cute.append(acc_shape, self.num_acc_stage)
)
if warp_idx == self.tma_warp_id:
ab_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_ab_stage
)
tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
tBgB1_slice = tBgB1[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
tBgB2_slice = tBgB2[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])]
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.mma_tiler[1] == 64):
slice_n = mma_tile_coord_mnl[1] // 2
tBgSFB1_slice = tBgSFB1[(None, slice_n, None, mma_tile_coord_mnl[2])]
tBgSFB2_slice = tBgSFB2[(None, slice_n, None, mma_tile_coord_mnl[2])]
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)
for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status)
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,
cache_policy=cutlass.Int64(cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()),
)
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,
cache_policy=cutlass.Int64(cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()),
)
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,
cache_policy=cutlass.Int64(cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()),
)
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,
cache_policy=cutlass.Int64(cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()),
)
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,
cache_policy=cutlass.Int64(cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()),
)
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,
cache_policy=cutlass.Int64(cutlass.Int64(_TMA_CACHE_EVICT_FIRST).ir_value()),
)
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)
ab_pipeline.producer_tail(ab_producer_state)
if warp_idx == self.mma_warp_id:
tmem.wait_for_alloc()
acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
sfa_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr + self.num_accumulator_tmem_cols,
dtype=self.sf_dtype,
)
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)
sfb1_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr + self.total_sfa_cols,
dtype=self.sf_dtype,
)
sfb2_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr + self.total_sfb_cols,
dtype=self.sf_dtype,
)
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)
tCtSFB2 = cute.make_tensor(sfb2_tmem_ptr, tCtSFB_layout)
copy_atom_s2t = cute.make_copy_atom(
tcgen05.Cp4x32x128bOp(self.cta_group),
self.sf_dtype,
)
(
tiled_copy_s2t_sfa,
tCsSFA_compact_s2t,
tCtSFA_compact_s2t,
) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA, copy_atom_s2t)
(
tiled_copy_s2t_sfb1,
tCsSFB1_compact_s2t,
tCtSFB1_compact_s2t,
) = self.mainloop_s2t_copy_and_partition(sSFB1, tCtSFB1, copy_atom_s2t)
(
tiled_copy_s2t_sfb2,
tCsSFB2_compact_s2t,
tCtSFB2_compact_s2t,
) = self.mainloop_s2t_copy_and_partition(sSFB2, tCtSFB2, copy_atom_s2t)
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
)
acc_stage_index = acc_producer_state.index
tCtAcc1 = tCtAcc_base[(None, None, None, acc_stage_index)]
tCtAcc2 = tCtAcc_base[(None, None, None, acc_stage_index + 1)]
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < k_tile_cnt:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
)
tCtSFB1_mma = tCtSFB1
tCtSFB2_mma = tCtSFB2
if cutlass.const_expr(self.mma_tiler[1] == 64):
offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
shifted_ptr1 = cute.recast_ptr(
acc_tmem_ptr + self.total_sfa_cols
+ offset,
dtype=self.sf_dtype,
)
tCtSFB1_mma = cute.make_tensor(shifted_ptr1, tCtSFB_layout)
shifted_ptr2 = cute.recast_ptr(
acc_tmem_ptr + self.total_sfb_cols
+ offset,
dtype=self.sf_dtype,
)
tCtSFB2_mma = cute.make_tensor(shifted_ptr2, tCtSFB_layout)
tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
for k_tile in range(k_tile_cnt):
ab_pipeline.consumer_wait(
ab_consumer_state, peek_ab_full_status
)
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,
)
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,
)
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,
)
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
ab_pipeline.consumer_release(ab_consumer_state)
ab_consumer_state.advance()
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < k_tile_cnt:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
)
acc_pipeline.producer_commit(acc_producer_state)
if warp_idx < self.mma_warp_id:
tmem.allocate(self.num_tmem_alloc_cols)
tmem.wait_for_alloc()
acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
copy_atom_t2r = sm100_utils.get_tmem_load_op(
self.mma_tiler,
self.c_layout,
self.c_dtype,
self.acc_dtype,
epi_tile,
False,
)
(
tiled_copy_t2r,
tTR_tAcc1_base,
tTR_rAcc1,
) = self.epilog_tmem_copy_and_partition(
tidx, tCtAcc_base, tCgC, epi_tile, copy_atom_t2r
)
tTR_tAcc2_base = tTR_tAcc1_base
tTR_rAcc2 = cute.make_rmem_tensor(tTR_rAcc1.shape, self.acc_dtype)
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_t2r, tTR_rC, tidx, sC
)
(
tma_atom_c,
bSG_sC,
bSG_gC,
) = self.epilog_gmem_copy_and_partition(
tma_atom_c, tCgC, epi_tile, sC, mma_tile_coord_mnl
)
acc_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, 1
)
tTR_tAcc1 = tTR_tAcc1_base[
(None, None, None, None, None, 0)
]
tTR_tAcc2 = tTR_tAcc2_base[
(None, None, None, None, None, 1)
]
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))
subtile_cnt = cute.size(tTR_tAcc1.shape, mode=[3])
for subtile_idx in cutlass.range(subtile_cnt):
tTR_tAcc_mn_1 = tTR_tAcc1[(None, None, None, subtile_idx)]
tTR_tAcc_mn_2 = tTR_tAcc2[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r, tTR_tAcc_mn_1, tTR_rAcc1)
cute.copy(tiled_copy_t2r, tTR_tAcc_mn_2, tTR_rAcc2)
acc1_vec = tTR_rAcc1.load()
acc2_vec = tTR_rAcc2.load()
acc2_vec = swiglu_f32(acc1_vec, acc2_vec)
acc2_vec = TensorSSA(acc2_vec, tRS_rC.layout, cutlass.Float16)
tRS_rC.store(acc2_vec)
cute.copy(
tiled_copy_r2s,
tRS_rC,
tRS_sC[(None, None, None, subtile_idx)],
)
if warp_idx == self.epilog_warp_id[0]:
cute.copy(
tma_atom_c,
bSG_sC[(None, subtile_idx)],
bSG_gC[(None, subtile_idx)],
)
self.epilog_sync_barrier.arrive_and_wait()
self.epilog_sync_barrier.arrive_and_wait()
self.epilog_sync_barrier.arrive_and_wait()
tmem.relinquish_alloc_permit()
tmem.free(acc_tmem_ptr)
def mainloop_s2t_copy_and_partition(
self,
sSF: cute.Tensor,
tSF: cute.Tensor,
copy_atom_s2t,
) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
tCsSF_compact = cute.filter_zeros(sSF)
tCtSF_compact = cute.filter_zeros(tSF)
tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact)
thr_copy_s2t = tiled_copy_s2t.get_slice(0)
tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
tiled_copy_s2t, tCsSF_compact_s2t_
)
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,
copy_atom_t2r,
) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
tAcc_epi = cute.flat_divide(
tAcc[((None, None), 0, 0, None)],
epi_tile,
)
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)
tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi)
gC_mnl_epi = cute.flat_divide(
gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
)
tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
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]:
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)
thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
tRS_sC = thr_copy_r2s.partition_D(sC)
tRS_rC = tiled_copy_r2s.retile(tTR_rC)
return tiled_copy_r2s, tRS_rC, tRS_sC
def epilog_gmem_copy_and_partition(
self,
atom: Union[cute.CopyAtom, cute.TiledCopy],
gC_mnl: cute.Tensor,
epi_tile: cute.Tile,
sC: cute.Tensor,
tile_coord: Tuple,
) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]:
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)
bSG_sC, bSG_gC_partitioned = cpasync.tma_partition(
tma_atom_c,
0,
cute.make_layout(1),
sC_for_tma_partition,
gC_for_tma_partition,
)
bSG_gC = bSG_gC_partitioned[(None, None, None, *tile_coord)]
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]:
num_acc_stage = 2
num_c_stage = 1
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 = ( #39kb
cute.size_in_bytes(a_dtype, a_smem_layout_stage_one)
+ cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) * 2
+ cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one)
+ cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) * 2
)
mbar_helpers_bytes = 128
c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) #8kb
c_bytes = c_bytes_per_stage * num_c_stage #8kb
num_ab_stage = (
smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes)
) // ab_bytes_per_stage
#232 - 128 - 8kb // 39 kb = 5, 4 for large shape
#c is smaller than ab, we might be able to fit more C stages at the tail of smem
num_c_stage += ( #4 for small shape, 1 for large shape
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
#-----------------------------------------------------------
_CACHE_SMALL, _CACHE_LARGE = None, None
def compile_kernel():
global _CACHE_SMALL, _CACHE_LARGE
if _CACHE_SMALL is not None:
return _CACHE_SMALL, _CACHE_LARGE
_max_active_clusters = 4096
gemm_small = Sm100BlockScaledSmall(
sf_vec_size, (256, 64), cluster_shape_mn=(2, 1), occupancy=OCCUPANCY
)
gemm_large = Sm100BlockScaledDenseGemmKernel(
sf_vec_size, (128, 128), cluster_shape_mn=(1, 1), occupancy=OCCUPANCY
)
# 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
)
_CACHE_SMALL = cute.compile(gemm_small, a_ptr,
b1_ptr,
b2_ptr,
sfa_ptr,
sfb1_ptr,
sfb2_ptr,
c_ptr,
(0, 0, 0, 0),
_max_active_clusters,
options="--opt-level 2 --gpu-arch sm_100a")
_CACHE_LARGE = cute.compile(gemm_large, a_ptr,
b1_ptr,
b2_ptr,
sfa_ptr,
sfb1_ptr,
sfb2_ptr,
c_ptr,
(0, 0, 0, 0),
options="--opt-level 2 --gpu-arch sm_100a")
return _CACHE_SMALL, _CACHE_LARGE
def compile_small_cluster():
gemm = Sm100BlockScaledDenseGemmKernel(
sf_vec_size, (128, 128), cluster_shape_mn=(1, 1), occupancy=OCCUPANCY
)
# 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
)
comp = cute.compile(gemm, a_ptr,
b1_ptr,
b2_ptr,
sfa_ptr,
sfb1_ptr,
sfb2_ptr,
c_ptr,
(0, 0, 0, 0),
options="--opt-level 2 --gpu-arch sm_100a")
return comp
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
_, k, _ = a.shape
m, n, l = c.shape
k = k * 2
small, large = compile_kernel()
if m <= 256:
compiled_func = small
elif n < 3000:
compiled_func = compile_small_cluster()
else:
compiled_func = large
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
)
compiled_func(a_ptr, b1_ptr, b2_ptr, sfa_ptr, sfb1_ptr, sfb2_ptr, c_ptr, (m, n, k, l))
return c
scrolls · 2613 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 382558.
⋯ 2487 unchanged linesc_ptr,(0, 0, 0, 0),_max_active_clusters,- options="--opt-level 3 --gpu-arch sm_100a")+ options="--opt-level 2 --gpu-arch sm_100a")_CACHE_LARGE = cute.compile(gemm_large, a_ptr,
Best evidence level for this revision: reported
JSON