submission 378438
novo_force · python · License unknown
Use it
Vendorable · source mirrored · license unknownView source →
No package. Vendor the mirrored source: 2790 lines, June 9 Researcher Reciprocity License v1.0.
submission.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-modal-nvfp4-dual-gemm-378438?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:acd3be5fe090d4ff26925f4d324872a18c20c36be485e2f0510fa74d27e363f3
license declaredunknown
license concludedunknown
authorsnovo_force
imported2026-08-15
Techniques
Extracted from the mirrored source by pattern, never inferred. Each row cites its line.
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.ONE,warp-specialization
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)Kernel source
submission.py2790 lines
from typing import Optional, Tuple, Type, Union
import os
_TARGET_ARCH = "sm_100a"
_SIGMOID_USE_EXP2_DEFAULT = True
_USE_RCP_DEFAULT = True
_SIGMOID_ENABLE_SAT = False
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
import torch
import torch.nn.functional as F
import cutlass
import cutlass.cute as cute
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.cute.math as cmath
from cutlass.cute.runtime import make_ptr
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
ab_dtype = cutlass.Float4E2M1FN
sf_dtype = cutlass.Float8E4M3FN
sf_vec_size = 16
_LOG2E = 1.4426950408889634
_SIGMOID_POS_TH = 8.0
_SIGMOID_NEG_TH = -16.0
_CUTE_EXP2 = getattr(cmath, "exp2", None)
_CUTE_RCP = getattr(cmath, "rcp", None)
if _CUTE_RCP is None:
_CUTE_RCP = getattr(cmath, "reciprocal", None)
_HAS_RCP = _CUTE_RCP is not None
def _sigmoid_exp_neg(x, use_exp2: bool):
if use_exp2 and _CUTE_EXP2 is not None:
return _CUTE_EXP2(x * cutlass.Float32(_LOG2E))
return cmath.exp(x, fastmath=True)
def _silu_mul_acc(acc1, acc2, one, zero, use_exp2: bool, use_rcp: bool):
if cutlass.const_expr(_SIGMOID_ENABLE_SAT):
pos = acc1 > cutlass.Float32(_SIGMOID_POS_TH)
neg = acc1 < cutlass.Float32(_SIGMOID_NEG_TH)
den = one + _sigmoid_exp_neg(zero - acc1, use_exp2)
if cutlass.const_expr(use_rcp and _HAS_RCP):
inv = _CUTE_RCP(den)
else:
inv = cmath.rsqrt(den, fastmath=True)
inv = inv * inv
inv = cute.where(pos, one, inv)
inv = cute.where(neg, zero, inv)
return (acc1 * acc2) * inv
den = one + _sigmoid_exp_neg(zero - acc1, use_exp2)
if cutlass.const_expr(use_rcp and _HAS_RCP):
inv = _CUTE_RCP(den)
else:
inv = cmath.rsqrt(den, fastmath=True)
inv = inv * inv
return (acc1 * acc2) * inv
class Sm100BlockScaledDenseGemmKernel:
def __init__(
self,
mma_tiler_mn: Tuple[int, int],
cluster_shape_mn: Tuple[int, int],
c_dtype: Type[cutlass.Numeric],
):
self.ab_dtype = cutlass.Float4E2M1FN
self.sf_dtype = cutlass.Float8E4M3FN
self.acc_dtype = cutlass.Float32
self.c_dtype = c_dtype
self.sf_vec_size = 16
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.mma_tiler_mn = mma_tiler_mn
self.cluster_shape_mn = cluster_shape_mn
self.occupancy = 1
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
self.num_tmem_alloc_cols = 512
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)),
)
def _setup_attributes(self):
tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.ab_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
tcgen05.CtaGroup.ONE,
self.mma_tiler_mn,
)
mma_inst_tile_k = 4
mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
self.mma_tiler = (
self.mma_tiler_mn[0],
self.mma_tiler_mn[1],
mma_inst_shape_k * mma_inst_tile_k,
)
self.cta_tile_shape_mnk = (
self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
self.mma_tiler[1],
self.mma_tiler[2],
)
self.cluster_layout_vmnk = cute.tiled_divide(
cute.make_layout((*self.cluster_shape_mn, 1)),
(tiled_mma.thr_id.shape,),
)
self.mma_inst_shape_mn_sfb = (
self.mma_tiler_mn[0],
cute.round_up(self.mma_tiler_mn[1], 128),
)
tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.ab_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
tcgen05.CtaGroup.ONE,
self.mma_inst_shape_mn_sfb,
)
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.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,
False,
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,
)
ab_stage_cap = 0
if self.mma_tiler_mn == (128, 128):
ab_stage_cap = 4
if ab_stage_cap > 0 and self.num_ab_stage > ab_stage_cap:
self.num_ab_stage = ab_stage_cap
self.prefetch_stage = self.num_ab_stage
self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
tiled_mma,
self.mma_tiler,
self.ab_dtype,
self.num_ab_stage,
)
self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
tiled_mma,
self.mma_tiler,
self.ab_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,
)
@cute.jit
def __call__(
self,
a_ptr: cute.Pointer,
b_ptr: cute.Pointer,
sfa_ptr: cute.Pointer,
sfb_ptr: cute.Pointer,
c_ptr: cute.Pointer,
m: cutlass.Int32,
n: cutlass.Int32,
k: cutlass.Int32,
l: cutlass.Int32,
):
self.a_dtype: Type[cutlass.Numeric] = a_ptr.value_type
self.b_dtype: Type[cutlass.Numeric] = b_ptr.value_type
self.sf_dtype: Type[cutlass.Numeric] = sfa_ptr.value_type
self.c_dtype: Type[cutlass.Numeric] = c_ptr.value_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()
a_tensor = cute.make_tensor(
a_ptr,
cute.make_ordered_layout(
(cute.assume(m, 32), k, l), order=(1, 0, 2)
),
)
b_tensor = cute.make_tensor(
b_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(
b_tensor.shape, self.sf_vec_size
)
sfb_tensor = cute.make_tensor(sfb_ptr, sfb_layout)
tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.ab_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
tcgen05.CtaGroup.ONE,
self.mma_tiler_mn,
)
tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.ab_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
tcgen05.CtaGroup.ONE,
self.mma_inst_shape_mn_sfb,
)
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_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
b_op,
b_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_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B(
sfb_op,
sfb_tensor,
sfb_smem_layout,
self.mma_tiler_sfb,
tiled_mma_sfb,
self.cluster_layout_sfb_vmnk.shape,
internal_type=cutlass.Int16,
)
atom_thr_size = cute.size(tiled_mma.thr_id.shape)
a_copy_size = cute.size_in_bytes(self.ab_dtype, a_smem_layout)
b_copy_size = cute.size_in_bytes(self.ab_dtype, b_smem_layout)
sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)
sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout)
self.num_tma_load_bytes = (
a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size
) * atom_thr_size
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]
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,
]
sB: 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,
]
sSFB: cute.struct.Align[
cute.struct.MemRange[
self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)
],
self.buffer_align_bytes,
]
self.shared_storage = SharedStorage
self.kernel(
tiled_mma,
tiled_mma_sfb,
tma_atom_a,
tma_tensor_a,
tma_atom_b,
tma_tensor_b,
tma_atom_sfa,
tma_tensor_sfa,
tma_atom_sfb,
tma_tensor_sfb,
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),
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_b: cute.CopyAtom,
mB_nkl: cute.Tensor,
tma_atom_sfa: cute.CopyAtom,
mSFA_mkl: cute.Tensor,
tma_atom_sfb: cute.CopyAtom,
mSFB_nkl: cute.Tensor,
tma_atom_c: Optional[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, None],
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_b)
cpasync.prefetch_descriptor(tma_atom_sfa)
cpasync.prefetch_descriptor(tma_atom_sfb)
cpasync.prefetch_descriptor(tma_atom_c)
bidx, bidy, bidz = cute.arch.block_idx()
mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
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)
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],
)
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.PipelineTmaAsync.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,
)
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
barrier_for_retrieve=self.tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
)
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)
sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner)
sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged)
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)
)
gB_nkl = cute.local_tile(
mB_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)
)
gSFB_nkl = cute.local_tile(
mSFB_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)
tCgB = thr_mma.partition_B(gB_nkl)
tCgSFA = thr_mma.partition_A(gSFA_mkl)
tCgSFB = thr_mma_sfb.partition_B(gSFB_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
)
tBsB, tBgB = cpasync.tma_partition(
tma_atom_b,
block_in_cluster_coord_vmnk[1],
b_cta_layout,
cute.group_modes(sB, 0, 3),
cute.group_modes(tCgB, 0, 3),
)
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
)
tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition(
tma_atom_sfb,
block_in_cluster_coord_sfb_vmnk[1],
sfb_cta_layout,
cute.group_modes(sSFB, 0, 3),
cute.group_modes(tCgSFB, 0, 3),
)
tBsSFB = cute.filter_zeros(tBsSFB)
tBgSFB = cute.filter_zeros(tBgSFB)
tCrA = tiled_mma.make_fragment_A(sA)
tCrB = tiled_mma.make_fragment_B(sB)
acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
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
)
tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])]
tBgB_slice = tBgB[(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
tBgSFB_slice = tBgSFB[(None, slice_n, None, mma_tile_coord_mnl[2])]
for prefetch_tile in cutlass.range(0, self.prefetch_stage, unroll=1):
cute.prefetch(tma_atom_a, tAgA_slice[(None, prefetch_tile)])
cute.prefetch(tma_atom_b, tBgB_slice[(None, prefetch_tile)])
cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, prefetch_tile)])
cute.prefetch(tma_atom_sfb, tBgSFB_slice[(None, prefetch_tile)])
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_b,
tBgB_slice[(None, ab_producer_state.count)],
tBsB[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=b_full_mcast_mask,
)
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_sfb,
tBgSFB_slice[(None, ab_producer_state.count)],
tBsSFB[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=sfb_full_mcast_mask,
)
if k_block_idx < k_block_cnt - self.prefetch_stage:
next_k_idx = ab_producer_state.count + self.prefetch_stage
cute.prefetch(tma_atom_a, tAgA_slice[(None, next_k_idx)])
cute.prefetch(tma_atom_b, tBgB_slice[(None, next_k_idx)])
cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, next_k_idx)])
cute.prefetch(tma_atom_sfb, tBgSFB_slice[(None, next_k_idx)])
ab_producer_state.advance()
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)
elif warp_idx == self.mma_warp_id:
tmem.wait_for_alloc()
acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
sfa_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc),
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)
sfb_tmem_ptr = cute.recast_ptr(
acc_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc)
+ 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)),
)
tCtSFB = cute.make_tensor(sfb_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_sfb, tCsSFB_compact_s2t, tCtSFB_compact_s2t = (
self.mainloop_s2t_copy_and_partition(sSFB, tCtSFB)
)
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
)
peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state)
tCtSFB_mma = tCtSFB
if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
shifted_ptr = cute.recast_ptr(
acc_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+ offset,
dtype=self.sf_dtype,
)
tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout)
for k_block_idx in range(k_block_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]
tCsSFB_compact_s2t_staged = tCsSFB_compact_s2t[s2t_stage_coord]
cute.copy(
tiled_copy_s2t_sfa,
tCsSFA_compact_s2t_staged,
tCtSFA_compact_s2t,
)
cute.copy(
tiled_copy_s2t_sfb,
tCsSFB_compact_s2t_staged,
tCtSFB_compact_s2t,
)
num_kphases = cute.size(tCrA, mode=[2])
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
kphase_coord = (None, None, kphase_idx, ab_consumer_state.index)
sf_kphase_coord = (None, None, kphase_idx)
tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kphase_coord].iterator)
tiled_mma.set(tcgen05.Field.SFB, tCtSFB_mma[sf_kphase_coord].iterator)
if k_block_idx == 0 and kphase_idx == 0:
tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
else:
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
cute.gemm(
tiled_mma,
tCtAcc,
tCrA[kphase_coord],
tCrB[kphase_coord],
tCtAcc,
)
ab_pipeline.consumer_release(ab_consumer_state)
ab_consumer_state.advance()
if ab_consumer_state.count < k_block_cnt:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
)
acc_pipeline.producer_commit(acc_producer_state)
elif warp_idx in self.epilog_warp_id:
tmem.allocate(self.num_tmem_alloc_cols)
tmem.wait_for_alloc()
acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
tCtAcc = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout)
tiled_copy_t2r, tTR_tAcc, tTR_rAcc = self.epilog_tmem_copy_and_partition(
tidx, tCtAcc, tCgC, epi_tile
)
tTR_rC = cute.make_fragment(tTR_rAcc.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(
tidx, tma_atom_c, tCgC, epi_tile, sC
)
bSG_gC = bSG_gC[(None, None, None, *mma_tile_coord_mnl)]
acc_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
)
acc_pipeline.consumer_wait(acc_consumer_state)
tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc))
bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3])
for subtile_idx in range(subtile_cnt):
tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc)
tRS_rC.store(tTR_rAcc.load().to(self.c_dtype))
cute.copy(
tiled_copy_r2s,
tRS_rC,
tRS_sC[(None, None, None, subtile_idx)],
)
cute.arch.fence_view_async_shared()
if warp_idx == self.epilog_warp_id[0]:
cute.copy(
tma_atom_c,
bSG_sC[(None, subtile_idx)],
bSG_gC[(None, subtile_idx)],
)
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(tcgen05.CtaGroup.ONE),
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,
) -> 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,
False,
)
tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0)], epi_tile)
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)
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_fragment(
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
num_c_stage = 3
a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(
tiled_mma,
mma_tiler_mnk,
a_dtype,
1,
)
b_smem_layout_staged_one = sm100_utils.make_smem_layout_b(
tiled_mma,
mma_tiler_mnk,
b_dtype,
1,
)
sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa(
tiled_mma,
mma_tiler_mnk,
sf_vec_size,
1,
)
sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb(
tiled_mma,
mma_tiler_mnk,
sf_vec_size,
1,
)
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
smem_budget = smem_capacity // max(1, int(occupancy))
budget_for_ab = smem_budget - (mbar_helpers_bytes + c_bytes)
num_ab_stage = budget_for_ab // ab_bytes_per_stage
if num_ab_stage < 1:
num_ab_stage = 1
leftover = smem_budget - ab_bytes_per_stage * num_ab_stage - (mbar_helpers_bytes + c_bytes)
if leftover > 0:
num_c_stage += leftover // (c_bytes_per_stage)
if num_c_stage < 1:
num_c_stage = 1
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]:
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 Sm100BlockScaledDualGemmKernel:
def __init__(
self,
mma_tiler_mn: Tuple[int, int],
cluster_shape_mn: Tuple[int, int],
c_dtype: Type[cutlass.Numeric],
ab_stage_cap: int = 0,
occupancy: int = 1,
force_tmem_512: bool = False,
store_pipe: bool = False,
sigmoid_use_exp2: bool = True,
sigmoid_use_rcp: bool = True,
):
self.ab_dtype = cutlass.Float4E2M1FN
self.sf_dtype = cutlass.Float8E4M3FN
self.acc_dtype = cutlass.Float32
self.c_dtype = c_dtype
self.sf_vec_size = 16
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.mma_tiler_mn = mma_tiler_mn
self.cluster_shape_mn = cluster_shape_mn
self.ab_stage_cap = ab_stage_cap
occ = int(occupancy)
self.occupancy = 1 if occ <= 0 else occ
self.force_tmem_512 = bool(force_tmem_512)
self.store_pipe = bool(store_pipe)
self.sigmoid_use_exp2 = bool(sigmoid_use_exp2)
self.sigmoid_use_rcp = bool(sigmoid_use_rcp)
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
self.num_tmem_alloc_cols = 512
self.tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=2,
num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
)
self.epilog_sync_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=32 * len(self.epilog_warp_id),
)
def _setup_attributes(self, ref_ptr: cute.Pointer):
tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.ab_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
tcgen05.CtaGroup.ONE,
self.mma_tiler_mn,
)
mma_inst_tile_k = 4
mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
self.mma_tiler = (
self.mma_tiler_mn[0],
self.mma_tiler_mn[1],
mma_inst_shape_k * mma_inst_tile_k,
)
self.cta_tile_shape_mnk = (
self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
self.mma_tiler[1],
self.mma_tiler[2],
)
self.cluster_layout_vmnk = cute.tiled_divide(
cute.make_layout((*self.cluster_shape_mn, 1)),
(tiled_mma.thr_id.shape,),
)
self.mma_inst_shape_mn_sfb = (
self.mma_tiler_mn[0],
cute.round_up(self.mma_tiler_mn[1], 128),
)
tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.ab_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
tcgen05.CtaGroup.ONE,
self.mma_inst_shape_mn_sfb,
)
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.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.epi_tile = sm100_utils.compute_epilogue_tile_shape(
self.cta_tile_shape_mnk,
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,
)
ab_stage_cap = int(self.ab_stage_cap)
if ab_stage_cap > 0 and self.num_ab_stage > ab_stage_cap:
self.num_ab_stage = ab_stage_cap
self.prefetch_stage = self.num_ab_stage
self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
tiled_mma,
self.mma_tiler,
self.ab_dtype,
self.num_ab_stage,
)
self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
tiled_mma,
self.mma_tiler,
self.ab_dtype,
self.num_ab_stage,
)
self.b2_smem_layout_staged = sm100_utils.make_smem_layout_b(
tiled_mma,
self.mma_tiler,
self.ab_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.sfb2_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,
)
acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape)
fake_base_ptr = cute.recast_ptr(ref_ptr, dtype=self.acc_dtype)
tCtAcc_0 = cute.make_tensor(fake_base_ptr, tCtAcc_fake.layout)
acc_cols_0 = tcgen05.find_tmem_tensor_col_offset(tCtAcc_0)
tCtAcc_1 = cute.make_tensor(fake_base_ptr + acc_cols_0, tCtAcc_fake.layout)
acc_cols_1 = tcgen05.find_tmem_tensor_col_offset(tCtAcc_1)
sf_base_ptr = fake_base_ptr + acc_cols_0 + acc_cols_1
sfa_tmem_ptr = cute.recast_ptr(sf_base_ptr, dtype=self.sf_dtype)
tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
tiled_mma,
self.mma_tiler,
self.sf_vec_size,
cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)),
)
tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
sfa_cols = tcgen05.find_tmem_tensor_col_offset(tCtSFA)
tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
tiled_mma,
self.mma_tiler,
self.sf_vec_size,
cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)),
)
tCtSFB1_0 = cute.make_tensor(
cute.recast_ptr(sf_base_ptr + sfa_cols, dtype=self.sf_dtype),
tCtSFB_layout,
)
sfb1_cols = tcgen05.find_tmem_tensor_col_offset(tCtSFB1_0)
tCtSFB2_0 = cute.make_tensor(
cute.recast_ptr(sf_base_ptr + sfa_cols + sfb1_cols, dtype=self.sf_dtype),
tCtSFB_layout,
)
sfb2_cols = tcgen05.find_tmem_tensor_col_offset(tCtSFB2_0)
total_cols = acc_cols_0 + acc_cols_1 + sfa_cols + sfb1_cols + sfb2_cols
if self.mma_tiler_mn == (128, 128):
total_cols = acc_cols_0 + acc_cols_1 + sfa_cols + sfb1_cols
if self.cta_tile_shape_mnk[1] == 64:
total_cols += 2
needed_cols = int(total_cols)
if needed_cols <= 32:
self.num_tmem_alloc_cols = 32
elif needed_cols <= 64:
self.num_tmem_alloc_cols = 64
elif needed_cols <= 128:
self.num_tmem_alloc_cols = 128
elif needed_cols <= 256:
self.num_tmem_alloc_cols = 256
elif needed_cols <= 512:
self.num_tmem_alloc_cols = 512
else:
raise ValueError(f"TMEM cols budget exceeded: needed_cols={needed_cols}")
if self.force_tmem_512 and self.mma_tiler_mn == (128, 128) and self.num_tmem_alloc_cols < 512:
self.num_tmem_alloc_cols = 512
@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.Int32,
n: cutlass.Int32,
k: cutlass.Int32,
l: cutlass.Int32,
):
self.a_dtype: Type[cutlass.Numeric] = a_ptr.value_type
self.b_dtype: Type[cutlass.Numeric] = b1_ptr.value_type
self.sf_dtype: Type[cutlass.Numeric] = sfa_ptr.value_type
self.c_dtype: Type[cutlass.Numeric] = c_ptr.value_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(c_ptr)
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)
tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.ab_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
tcgen05.CtaGroup.ONE,
self.mma_tiler_mn,
)
tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.ab_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
tcgen05.CtaGroup.ONE,
self.mma_inst_shape_mn_sfb,
)
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
)
b1_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
b2_smem_layout = cute.slice_(self.b2_smem_layout_staged, (None, None, None, 0))
tma_atom_b1, tma_tensor_b1 = cute.nvgpu.make_tiled_tma_atom_B(
b_op,
b1_tensor,
b1_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,
b2_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
)
sfb1_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0))
sfb2_smem_layout = cute.slice_(self.sfb2_smem_layout_staged, (None, None, None, 0))
tma_atom_sfb1, tma_tensor_sfb1 = cute.nvgpu.make_tiled_tma_atom_B(
sfb_op,
sfb1_tensor,
sfb1_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,
sfb2_smem_layout,
self.mma_tiler_sfb,
tiled_mma_sfb,
self.cluster_layout_sfb_vmnk.shape,
internal_type=cutlass.Int16,
)
atom_thr_size = cute.size(tiled_mma.thr_id.shape)
a_copy_size = cute.size_in_bytes(self.ab_dtype, a_smem_layout)
b1_copy_size = cute.size_in_bytes(self.ab_dtype, b1_smem_layout)
b2_copy_size = cute.size_in_bytes(self.ab_dtype, b2_smem_layout)
sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)
sfb1_copy_size = cute.size_in_bytes(self.sf_dtype, sfb1_smem_layout)
sfb2_copy_size = cute.size_in_bytes(self.sf_dtype, sfb2_smem_layout)
self.num_tma_load_bytes = (
a_copy_size
+ b1_copy_size
+ b2_copy_size
+ sfa_copy_size
+ sfb1_copy_size
+ sfb2_copy_size
) * 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 = Sm100BlockScaledDenseGemmKernel._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]
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.b2_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.sfb2_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.b2_smem_layout_staged,
self.sfa_smem_layout_staged,
self.sfb_smem_layout_staged,
self.sfb2_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),
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: Optional[cute.CopyAtom],
mC_mnl: cute.Tensor,
cluster_layout_vmnk: cute.Layout,
cluster_layout_sfb_vmnk: cute.Layout,
a_smem_layout_staged: cute.ComposedLayout,
b1_smem_layout_staged: cute.ComposedLayout,
b2_smem_layout_staged: cute.ComposedLayout,
sfa_smem_layout_staged: cute.Layout,
sfb1_smem_layout_staged: cute.Layout,
sfb2_smem_layout_staged: cute.Layout,
c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None],
epi_tile: cute.Tile,
):
warp_idx = cute.arch.make_warp_uniform(cute.arch.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_v = bidx % cute.size(tiled_mma.thr_id.shape)
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)
mma_tile_coord_mnl = (
bidx // cute.size(tiled_mma.thr_id.shape),
bidy,
bidz,
)
tidx, _, _ = cute.arch.thread_idx()
smem = utils.SmemAllocator()
storage = smem.allocate(self.shared_storage)
ab_pipeline = pipeline.PipelineTmaAsync.create(
barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
num_stages=self.num_ab_stage,
producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
consumer_group=pipeline.CooperativeGroup(
pipeline.Agent.Thread, self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
),
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
acc_pipeline = pipeline.PipelineUmmaAsync.create(
barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
num_stages=self.num_acc_stage,
producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread),
consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, len(self.epilog_warp_id)),
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],
)
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(
b1_smem_layout_staged.outer, swizzle=b1_smem_layout_staged.inner
)
sB2 = storage.sB2.get_tensor(
b2_smem_layout_staged.outer, swizzle=b2_smem_layout_staged.inner
)
sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
sSFB1 = storage.sSFB1.get_tensor(sfb1_smem_layout_staged)
sSFB2 = storage.sSFB2.get_tensor(sfb2_smem_layout_staged)
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),
)
tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition(
tma_atom_sfa,
block_in_cluster_coord_vmnk[2],
a_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),
)
tBsSFB1 = cute.filter_zeros(tBsSFB1)
tBgSFB1 = cute.filter_zeros(tBgSFB1)
tBsSFB2 = cute.filter_zeros(tBsSFB2)
tBgSFB2 = cute.filter_zeros(tBgSFB2)
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(acc_shape)
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
)
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])]
for prefetch_tile in cutlass.range(0, self.prefetch_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_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)
bar_ptr = ab_pipeline.producer_get_barrier(ab_producer_state)
idx = ab_producer_state.index
cnt = ab_producer_state.count
cute.copy(
tma_atom_a,
tAgA_slice[(None, cnt)],
tAsA[(None, idx)],
tma_bar_ptr=bar_ptr,
mcast_mask=a_full_mcast_mask,
)
cute.copy(
tma_atom_b1,
tBgB1_slice[(None, cnt)],
tBsB1[(None, idx)],
tma_bar_ptr=bar_ptr,
mcast_mask=b_full_mcast_mask,
)
cute.copy(
tma_atom_b2,
tBgB2_slice[(None, cnt)],
tBsB2[(None, idx)],
tma_bar_ptr=bar_ptr,
mcast_mask=b_full_mcast_mask,
)
cute.copy(
tma_atom_sfa,
tAgSFA_slice[(None, cnt)],
tAsSFA[(None, idx)],
tma_bar_ptr=bar_ptr,
mcast_mask=sfa_full_mcast_mask,
)
cute.copy(
tma_atom_sfb1,
tBgSFB1_slice[(None, cnt)],
tBsSFB1[(None, idx)],
tma_bar_ptr=bar_ptr,
mcast_mask=sfb_full_mcast_mask,
)
cute.copy(
tma_atom_sfb2,
tBgSFB2_slice[(None, cnt)],
tBsSFB2[(None, idx)],
tma_bar_ptr=bar_ptr,
mcast_mask=sfb_full_mcast_mask,
)
if k_block_idx < k_block_cnt - self.prefetch_stage:
next_k_idx = cnt + self.prefetch_stage
cute.prefetch(tma_atom_a, tAgA_slice[(None, next_k_idx)])
cute.prefetch(tma_atom_b1, tBgB1_slice[(None, next_k_idx)])
cute.prefetch(tma_atom_b2, tBgB2_slice[(None, next_k_idx)])
cute.prefetch(tma_atom_sfa, tAgSFA_slice[(None, next_k_idx)])
cute.prefetch(tma_atom_sfb1, tBgSFB1_slice[(None, next_k_idx)])
cute.prefetch(tma_atom_sfb2, tBgSFB2_slice[(None, next_k_idx)])
ab_producer_state.advance()
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)
elif warp_idx == self.mma_warp_id:
tmem.wait_for_alloc()
base_ptr = tmem.retrieve_ptr(self.acc_dtype)
tCtAcc1 = cute.make_tensor(base_ptr, tCtAcc_fake.layout)
base_ptr_2 = base_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc1)
tCtAcc2 = cute.make_tensor(base_ptr_2, tCtAcc_fake.layout)
sf_base_ptr = base_ptr_2 + tcgen05.find_tmem_tensor_col_offset(tCtAcc2)
sfa_tmem_ptr = cute.recast_ptr(sf_base_ptr, 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(
sf_base_ptr + 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_(sfb1_smem_layout_staged, (None, None, None, 0)),
)
tCtSFB1 = cute.make_tensor(sfb1_tmem_ptr, tCtSFB_layout)
tCtSFB2 = tCtSFB1
if cutlass.const_expr(self.mma_tiler_mn != (128, 128)):
sfb2_tmem_ptr = cute.recast_ptr(
sf_base_ptr
+ 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 = (
Sm100BlockScaledDenseGemmKernel.mainloop_s2t_copy_and_partition(self, sSFA, tCtSFA)
)
tiled_copy_s2t_sfb1, tCsSFB1_compact_s2t, tCtSFB1_compact_s2t = (
Sm100BlockScaledDenseGemmKernel.mainloop_s2t_copy_and_partition(self, sSFB1, tCtSFB1)
)
tiled_copy_s2t_sfb2, tCsSFB2_compact_s2t, tCtSFB2_compact_s2t = (
Sm100BlockScaledDenseGemmKernel.mainloop_s2t_copy_and_partition(self, 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
)
peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state)
tCtSFB1_mma = tCtSFB1
tCtSFB2_mma = tCtSFB2
if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
shift_base = sf_base_ptr + tcgen05.find_tmem_tensor_col_offset(tCtSFA) + offset
tCtSFB1_mma = cute.make_tensor(
cute.recast_ptr(shift_base, dtype=self.sf_dtype), tCtSFB_layout
)
shift_base_2 = (
sf_base_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFB1)
+ offset
)
tCtSFB2_mma = cute.make_tensor(
cute.recast_ptr(shift_base_2, dtype=self.sf_dtype), tCtSFB_layout
)
for k_block_idx in range(k_block_cnt):
ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status)
s2t_stage_coord = (None, None, None, None, ab_consumer_state.index)
cute.copy(
tiled_copy_s2t_sfa,
tCsSFA_compact_s2t[s2t_stage_coord],
tCtSFA_compact_s2t,
)
num_kphases = cute.size(tCrA, mode=[2])
if cutlass.const_expr(self.mma_tiler_mn == (128, 128)):
cute.copy(
tiled_copy_s2t_sfb1,
tCsSFB1_compact_s2t[s2t_stage_coord],
tCtSFB1_compact_s2t,
)
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
kphase_coord = (None, None, kphase_idx, ab_consumer_state.index)
sf_kphase_coord = (None, None, kphase_idx)
tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kphase_coord].iterator)
first = k_block_idx == 0 and kphase_idx == 0
tiled_mma.set(
tcgen05.Field.SFB, tCtSFB1_mma[sf_kphase_coord].iterator
)
tiled_mma.set(tcgen05.Field.ACCUMULATE, not first)
cute.gemm(
tiled_mma,
tCtAcc1,
tCrA[kphase_coord],
tCrB1[kphase_coord],
tCtAcc1,
)
cute.copy(
tiled_copy_s2t_sfb2,
tCsSFB2_compact_s2t[s2t_stage_coord],
tCtSFB2_compact_s2t,
)
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
kphase_coord = (None, None, kphase_idx, ab_consumer_state.index)
sf_kphase_coord = (None, None, kphase_idx)
tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kphase_coord].iterator)
first = k_block_idx == 0 and kphase_idx == 0
tiled_mma.set(
tcgen05.Field.SFB, tCtSFB1_mma[sf_kphase_coord].iterator
)
tiled_mma.set(tcgen05.Field.ACCUMULATE, not first)
cute.gemm(
tiled_mma,
tCtAcc2,
tCrA[kphase_coord],
tCrB2[kphase_coord],
tCtAcc2,
)
else:
cute.copy(
tiled_copy_s2t_sfb1,
tCsSFB1_compact_s2t[s2t_stage_coord],
tCtSFB1_compact_s2t,
)
cute.copy(
tiled_copy_s2t_sfb2,
tCsSFB2_compact_s2t[s2t_stage_coord],
tCtSFB2_compact_s2t,
)
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
kphase_coord = (None, None, kphase_idx, ab_consumer_state.index)
sf_kphase_coord = (None, None, kphase_idx)
tiled_mma.set(tcgen05.Field.SFA, tCtSFA[sf_kphase_coord].iterator)
first = k_block_idx == 0 and kphase_idx == 0
tiled_mma.set(
tcgen05.Field.SFB, tCtSFB1_mma[sf_kphase_coord].iterator
)
tiled_mma.set(tcgen05.Field.ACCUMULATE, not first)
cute.gemm(
tiled_mma,
tCtAcc1,
tCrA[kphase_coord],
tCrB1[kphase_coord],
tCtAcc1,
)
tiled_mma.set(
tcgen05.Field.SFB, tCtSFB2_mma[sf_kphase_coord].iterator
)
tiled_mma.set(tcgen05.Field.ACCUMULATE, not first)
cute.gemm(
tiled_mma,
tCtAcc2,
tCrA[kphase_coord],
tCrB2[kphase_coord],
tCtAcc2,
)
ab_pipeline.consumer_release(ab_consumer_state)
ab_consumer_state.advance()
if ab_consumer_state.count < k_block_cnt:
peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state)
acc_pipeline.producer_commit(acc_producer_state)
elif warp_idx in self.epilog_warp_id:
tmem.allocate(self.num_tmem_alloc_cols)
tmem.wait_for_alloc()
base_ptr = tmem.retrieve_ptr(self.acc_dtype)
tCtAcc1 = cute.make_tensor(base_ptr, tCtAcc_fake.layout)
base_ptr_2 = base_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc1)
tCtAcc2 = cute.make_tensor(base_ptr_2, tCtAcc_fake.layout)
acc_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
)
acc_pipeline.consumer_wait(acc_consumer_state)
tiled_copy_t2r_1, tTR_tAcc_1, tTR_rAcc_1 = Sm100BlockScaledDenseGemmKernel.epilog_tmem_copy_and_partition(
self, tidx, tCtAcc1, tCgC, epi_tile
)
tTR_rC = cute.make_fragment(tTR_rAcc_1.shape, self.c_dtype)
tiled_copy_r2s, tRS_rC, tRS_sC = Sm100BlockScaledDenseGemmKernel.epilog_smem_copy_and_partition(
self, tiled_copy_t2r_1, tTR_rC, tidx, sC
)
tma_atom_c, bSG_sC, bSG_gC = Sm100BlockScaledDenseGemmKernel.epilog_gmem_copy_and_partition(
self, tidx, tma_atom_c, tCgC, epi_tile, sC
)
bSG_gC = bSG_gC[(None, None, None, *mma_tile_coord_mnl)]
tTR_tAcc_1 = cute.group_modes(tTR_tAcc_1, 3, cute.rank(tTR_tAcc_1))
bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
tiled_copy_t2r_2, tTR_tAcc_2, tTR_rAcc_2 = Sm100BlockScaledDenseGemmKernel.epilog_tmem_copy_and_partition(
self, tidx, tCtAcc2, tCgC, epi_tile
)
tTR_tAcc_2 = cute.group_modes(tTR_tAcc_2, 3, cute.rank(tTR_tAcc_2))
subtile_cnt = cute.size(tTR_tAcc_1.shape, mode=[3])
one = cutlass.Float32(1.0)
zero = cutlass.Float32(0.0)
batch = self.num_c_stage
tail = subtile_cnt % batch
main_cnt = subtile_cnt - tail
c_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, 32 * len(self.epilog_warp_id)
)
c_pipeline = pipeline.PipelineTmaStore.create(
num_stages=self.num_c_stage,
producer_group=c_producer_group,
)
if cutlass.const_expr(self.store_pipe or self.cta_tile_shape_mnk[1] == 128):
if warp_idx == self.epilog_warp_id[0]:
c_pipeline.producer_acquire()
self.epilog_sync_barrier.sync()
if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
for base in range(0, main_cnt, batch):
for off in range(batch):
subtile_idx = base + off
c_buf = subtile_idx % self.num_c_stage
if base != 0 and off == 0 and warp_idx != self.epilog_warp_id[0]:
cute.copy(
tiled_copy_r2s,
tRS_rC,
tRS_sC[(None, None, None, c_buf)],
)
else:
tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)
acc1 = tTR_rAcc_1.load()
tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)
acc2 = tTR_rAcc_2.load()
out = _silu_mul_acc(
acc1,
acc2,
one,
zero,
self.sigmoid_use_exp2,
self.sigmoid_use_rcp,
)
tRS_rC.store(out.to(self.c_dtype))
cute.copy(
tiled_copy_r2s,
tRS_rC,
tRS_sC[(None, None, None, c_buf)],
)
cute.arch.fence_view_async_shared()
self.epilog_sync_barrier.sync()
if base + batch < main_cnt and warp_idx != self.epilog_warp_id[0]:
subtile_idx = base + batch
tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)
acc1 = tTR_rAcc_1.load()
tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)
acc2 = tTR_rAcc_2.load()
out = _silu_mul_acc(
acc1,
acc2,
one,
zero,
self.sigmoid_use_exp2,
self.sigmoid_use_rcp,
)
tRS_rC.store(out.to(self.c_dtype))
if warp_idx == self.epilog_warp_id[0]:
for off in range(batch):
subtile_idx = base + off
c_buf = subtile_idx % self.num_c_stage
cute.copy(
tma_atom_c,
bSG_sC[(None, c_buf)],
bSG_gC[(None, subtile_idx)],
)
c_pipeline.producer_commit()
c_pipeline.producer_acquire()
self.epilog_sync_barrier.sync()
if tail != 0:
base = main_cnt
for off in range(tail):
subtile_idx = base + off
c_buf = subtile_idx % self.num_c_stage
tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)
acc1 = tTR_rAcc_1.load()
tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)
acc2 = tTR_rAcc_2.load()
out = _silu_mul_acc(
acc1,
acc2,
one,
zero,
self.sigmoid_use_exp2,
self.sigmoid_use_rcp,
)
tRS_rC.store(out.to(self.c_dtype))
cute.copy(
tiled_copy_r2s,
tRS_rC,
tRS_sC[(None, None, None, c_buf)],
)
cute.arch.fence_view_async_shared()
self.epilog_sync_barrier.sync()
if warp_idx == self.epilog_warp_id[0]:
for off in range(tail):
subtile_idx = base + off
c_buf = subtile_idx % self.num_c_stage
cute.copy(
tma_atom_c,
bSG_sC[(None, c_buf)],
bSG_gC[(None, subtile_idx)],
)
c_pipeline.producer_commit()
c_pipeline.producer_acquire()
self.epilog_sync_barrier.sync()
self.epilog_sync_barrier.sync()
if warp_idx == self.epilog_warp_id[0]:
c_pipeline.producer_tail()
self.epilog_sync_barrier.sync()
else:
for base in range(0, main_cnt, batch):
for off in range(batch):
subtile_idx = base + off
tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)
acc1 = tTR_rAcc_1.load()
tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)
acc2 = tTR_rAcc_2.load()
out = _silu_mul_acc(
acc1,
acc2,
one,
zero,
self.sigmoid_use_exp2,
self.sigmoid_use_rcp,
)
tRS_rC.store(out.to(self.c_dtype))
cute.copy(
tiled_copy_r2s,
tRS_rC,
tRS_sC[(None, None, None, subtile_idx)],
)
cute.arch.fence_view_async_shared()
self.epilog_sync_barrier.sync()
if warp_idx == self.epilog_warp_id[0]:
for off in range(batch):
store_idx = base + off
cute.copy(
tma_atom_c,
bSG_sC[(None, store_idx)],
bSG_gC[(None, store_idx)],
)
c_pipeline.producer_commit()
c_pipeline.producer_acquire()
self.epilog_sync_barrier.sync()
if tail != 0:
base = main_cnt
for off in range(tail):
subtile_idx = base + off
tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)
acc1 = tTR_rAcc_1.load()
tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, subtile_idx)]
cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)
acc2 = tTR_rAcc_2.load()
out = _silu_mul_acc(
acc1,
acc2,
one,
zero,
self.sigmoid_use_exp2,
self.sigmoid_use_rcp,
)
tRS_rC.store(out.to(self.c_dtype))
cute.copy(
tiled_copy_r2s,
tRS_rC,
tRS_sC[(None, None, None, subtile_idx)],
)
cute.arch.fence_view_async_shared()
self.epilog_sync_barrier.sync()
if warp_idx == self.epilog_warp_id[0]:
for off in range(tail):
store_idx = base + off
cute.copy(
tma_atom_c,
bSG_sC[(None, store_idx)],
bSG_gC[(None, store_idx)],
)
c_pipeline.producer_commit()
c_pipeline.producer_acquire()
self.epilog_sync_barrier.sync()
self.epilog_sync_barrier.sync()
if warp_idx == self.epilog_warp_id[0]:
c_pipeline.producer_tail()
self.epilog_sync_barrier.sync()
else:
tTR_rAcc_1_b = cute.make_fragment(tTR_rAcc_1.shape, self.acc_dtype)
tTR_rAcc_2_b = cute.make_fragment(tTR_rAcc_2.shape, self.acc_dtype)
if subtile_cnt > 0:
tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, 0)]
cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)
tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, 0)]
cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)
for subtile_idx in range(subtile_cnt):
cur_is_0 = (subtile_idx & 1) == 0
next_idx = subtile_idx + 1
if next_idx < subtile_cnt:
if cur_is_0:
tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, next_idx)]
cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1_b)
tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, next_idx)]
cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2_b)
else:
tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, next_idx)]
cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)
tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, next_idx)]
cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)
if cur_is_0:
acc1 = tTR_rAcc_1.load()
acc2 = tTR_rAcc_2.load()
out = _silu_mul_acc(
acc1,
acc2,
one,
zero,
self.sigmoid_use_exp2,
self.sigmoid_use_rcp,
)
tRS_rC.store(out.to(self.c_dtype))
else:
acc1 = tTR_rAcc_1_b.load()
acc2 = tTR_rAcc_2_b.load()
out = _silu_mul_acc(
acc1,
acc2,
one,
zero,
self.sigmoid_use_exp2,
self.sigmoid_use_rcp,
)
tRS_rC.store(out.to(self.c_dtype))
cute.copy(
tiled_copy_r2s,
tRS_rC,
tRS_sC[(None, None, None, subtile_idx)],
)
if (subtile_idx + 1) % batch == 0:
cute.arch.fence_view_async_shared()
self.epilog_sync_barrier.sync()
if warp_idx == self.epilog_warp_id[0]:
for off in range(batch):
store_idx = subtile_idx - (batch - 1 - off)
cute.copy(
tma_atom_c,
bSG_sC[(None, store_idx)],
bSG_gC[(None, store_idx)],
)
if tail != 0:
cute.arch.fence_view_async_shared()
self.epilog_sync_barrier.sync()
if warp_idx == self.epilog_warp_id[0]:
base = subtile_cnt - tail
for off in range(batch):
if off < tail:
store_idx = base + off
cute.copy(
tma_atom_c,
bSG_sC[(None, store_idx)],
bSG_gC[(None, store_idx)],
)
tmem.relinquish_alloc_permit()
tmem.free(base_ptr)
@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
num_c_stage = 3
a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(
tiled_mma,
mma_tiler_mnk,
a_dtype,
1,
)
b_smem_layout_staged_one = sm100_utils.make_smem_layout_b(
tiled_mma,
mma_tiler_mnk,
b_dtype,
1,
)
sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa(
tiled_mma,
mma_tiler_mnk,
sf_vec_size,
1,
)
sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb(
tiled_mma,
mma_tiler_mnk,
sf_vec_size,
1,
)
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)
+ 2 * cute.size_in_bytes(b_dtype, b_smem_layout_staged_one)
+ cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one)
+ 2 * 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
smem_budget = smem_capacity // max(1, int(occupancy))
budget_for_ab = smem_budget - (mbar_helpers_bytes + c_bytes)
num_ab_stage = budget_for_ab // ab_bytes_per_stage
if num_ab_stage < 1:
num_ab_stage = 1
leftover = smem_budget - ab_bytes_per_stage * num_ab_stage - (mbar_helpers_bytes + c_bytes)
if leftover > 0:
num_c_stage += leftover // (c_bytes_per_stage)
if num_c_stage < 1:
num_c_stage = 1
return num_acc_stage, num_ab_stage, num_c_stage
_GEMM_KERNEL_CACHE = {}
_DUAL_GEMM_KERNEL_CACHE = {}
_CUDA_CACHE_CLEARED = False
_BUF_CACHE = {}
def _maybe_empty_cuda_cache_once() -> None:
global _CUDA_CACHE_CLEARED
if not _CUDA_CACHE_CLEARED:
torch.cuda.empty_cache()
_CUDA_CACHE_CLEARED = True
def _get_buf(tag: str, shape, stride, device, dtype: torch.dtype) -> torch.Tensor:
if isinstance(shape, torch.Size):
shape = tuple(int(x) for x in shape)
else:
shape = tuple(int(x) for x in shape)
if isinstance(stride, torch.Size):
stride = tuple(int(x) for x in stride)
else:
stride = tuple(int(x) for x in stride)
if (
len(shape) == 3
and shape[2] == 1
and len(stride) == 3
and stride[1] == 1
and stride[2] == 1
and stride[0] >= shape[1]
):
m, n, l = shape
key = (tag, int(stride[0]), int(n), device, dtype)
buf = _BUF_CACHE.get(key)
if buf is None or buf.shape[0] < m:
buf = torch.empty_strided((m, n, l), stride, device=device, dtype=dtype)
_BUF_CACHE[key] = buf
return buf[:m, :n, :l]
key = (tag, shape, stride, device, dtype)
buf = _BUF_CACHE.get(key)
if buf is None:
buf = torch.empty_strided(shape, stride, device=device, dtype=dtype)
_BUF_CACHE[key] = buf
return buf
def _compile_gemm_kernel(out_dtype: Type[cutlass.Numeric], mma_tiler_mn, cluster_shape_mn):
key = (out_dtype, mma_tiler_mn, cluster_shape_mn)
cached = _GEMM_KERNEL_CACHE.get(key)
if cached is not None:
return cached
target_arch = _TARGET_ARCH
cutlass.cuda.initialize_cuda_context()
a_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
b_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
c_ptr = make_ptr(out_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)
sfa_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)
sfb_ptr = make_ptr(sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32)
kernel = Sm100BlockScaledDenseGemmKernel(mma_tiler_mn, cluster_shape_mn, c_dtype=out_dtype)
compiled = cute.compile[cute.GPUArch(target_arch)](
kernel, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, 0, 0, 0, 0
)
_GEMM_KERNEL_CACHE[key] = compiled
_maybe_empty_cuda_cache_once()
return compiled
def _compile_dual_gemm_kernel(
out_dtype: Type[cutlass.Numeric],
mma_tiler_mn,
cluster_shape_mn,
ab_stage_cap: int,
occupancy: int,
force_tmem_512: bool,
store_pipe: bool,
sigmoid_use_exp2: bool,
sigmoid_use_rcp: bool,
):
key = (
out_dtype,
mma_tiler_mn,
cluster_shape_mn,
int(ab_stage_cap),
int(occupancy),
bool(force_tmem_512),
bool(store_pipe),
bool(sigmoid_use_exp2),
bool(sigmoid_use_rcp),
bool(_SIGMOID_ENABLE_SAT),
)
cached = _DUAL_GEMM_KERNEL_CACHE.get(key)
if cached is not None:
return cached
target_arch = _TARGET_ARCH
cutlass.cuda.initialize_cuda_context()
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(out_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)
kernel = Sm100BlockScaledDualGemmKernel(
mma_tiler_mn,
cluster_shape_mn,
c_dtype=out_dtype,
ab_stage_cap=int(ab_stage_cap),
occupancy=int(occupancy),
force_tmem_512=bool(force_tmem_512),
store_pipe=bool(store_pipe),
sigmoid_use_exp2=bool(sigmoid_use_exp2),
sigmoid_use_rcp=bool(sigmoid_use_rcp),
)
compiled = cute.compile[cute.GPUArch(target_arch)](
kernel,
a_ptr,
b1_ptr,
b2_ptr,
sfa_ptr,
sfb1_ptr,
sfb2_ptr,
c_ptr,
0,
0,
0,
0,
)
_DUAL_GEMM_KERNEL_CACHE[key] = compiled
_maybe_empty_cuda_cache_once()
return compiled
def _run_gemm(compiled, out_dtype, a, b, sfa_p, sfb_p, out):
m, k_half, l = a.shape
n, _, _ = b.shape
k = k_half * 2
compiled(
make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
make_ptr(ab_dtype, b.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
make_ptr(sf_dtype, sfa_p.data_ptr(), cute.AddressSpace.gmem, assumed_align=32),
make_ptr(sf_dtype, sfb_p.data_ptr(), cute.AddressSpace.gmem, assumed_align=32),
make_ptr(out_dtype, out.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
m,
n,
k,
l,
)
def _run_dual_gemm(compiled, out_dtype, a, b1, b2, sfa_p, sfb1_p, sfb2_p, out):
m, k_half, l = a.shape
n, _, _ = b1.shape
k = k_half * 2
compiled(
make_ptr(ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
make_ptr(ab_dtype, b1.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
make_ptr(ab_dtype, b2.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
make_ptr(sf_dtype, sfa_p.data_ptr(), cute.AddressSpace.gmem, assumed_align=32),
make_ptr(sf_dtype, sfb1_p.data_ptr(), cute.AddressSpace.gmem, assumed_align=32),
make_ptr(sf_dtype, sfb2_p.data_ptr(), cute.AddressSpace.gmem, assumed_align=32),
make_ptr(out_dtype, out.data_ptr(), cute.AddressSpace.gmem, assumed_align=16),
m,
n,
k,
l,
)
def _is_rank_shape(m: int, n: int, k: int, l: int) -> bool:
if l != 1:
return False
if m == 256 and n == 4096 and k == 7168:
return True
if m == 512 and n == 4096 and k == 7168:
return True
if m == 256 and n == 3072 and k == 4096:
return True
if m == 512 and n == 3072 and k == 7168:
return True
return False
_RANKED_CLUSTER_SHAPE_V0 = {
(256, 4096, 7168): (2, 2),
(512, 4096, 7168): (2, 2),
(256, 3072, 4096): (2, 2),
(512, 3072, 7168): (2, 2),
}
_RANKED_CLUSTER_SHAPE_V1 = {
(256, 4096, 7168): (2, 2),
(512, 4096, 7168): (4, 1),
(256, 3072, 4096): (2, 2),
(512, 3072, 7168): (4, 1),
}
_RANKED_DUAL_VARIANT = {
(256, 4096, 7168): 0,
(512, 4096, 7168): 1,
(256, 3072, 4096): 0,
(512, 3072, 7168): 1,
}
_RANKED_MMA_TILER_MN = {
(256, 4096, 7168): (128, 64),
(512, 4096, 7168): (128, 128),
(256, 3072, 4096): (128, 64),
(512, 3072, 7168): (128, 128),
}
_RANKED_AB_STAGE_CAP = {
(256, 4096, 7168): 99,
(512, 4096, 7168): 99,
(256, 3072, 4096): 99,
(512, 3072, 7168): 99,
}
_RANKED_OCCUPANCY = {
(256, 4096, 7168): 1,
(512, 4096, 7168): 1,
(256, 3072, 4096): 1,
(512, 3072, 7168): 1,
}
_RANKED_STORE_PIPE = {
(256, 4096, 7168): True,
(512, 4096, 7168): False,
(256, 3072, 4096): False,
(512, 3072, 7168): False,
}
_RANKED_FORCE_TMEM_512 = {
(256, 4096, 7168): False,
(512, 4096, 7168): False,
(256, 3072, 4096): False,
(512, 3072, 7168): True,
}
_RANKED_SIGMOID_EXP2 = {
(256, 4096, 7168): _SIGMOID_USE_EXP2_DEFAULT,
(512, 4096, 7168): _SIGMOID_USE_EXP2_DEFAULT,
(256, 3072, 4096): _SIGMOID_USE_EXP2_DEFAULT,
(512, 3072, 7168): False,
}
_RANKED_SIGMOID_RCP = {
(256, 4096, 7168): _USE_RCP_DEFAULT,
(512, 4096, 7168): _USE_RCP_DEFAULT,
(256, 3072, 4096): _USE_RCP_DEFAULT,
(512, 3072, 7168): True,
}
def _get_rank_dual_kernel(m: int, n: int, k: int):
shape = (m, n, k)
variant = int(_RANKED_DUAL_VARIANT.get(shape, 0))
mma_tiler_mn = _RANKED_MMA_TILER_MN.get(shape, (128, 64))
cluster_shape_mn = (_RANKED_CLUSTER_SHAPE_V1 if variant == 1 else _RANKED_CLUSTER_SHAPE_V0).get(shape, (2, 2))
ab_stage_cap = int(_RANKED_AB_STAGE_CAP.get(shape, 0))
occupancy = int(_RANKED_OCCUPANCY.get(shape, 1))
store_pipe = bool(_RANKED_STORE_PIPE.get(shape, False))
force_tmem_512 = bool(_RANKED_FORCE_TMEM_512.get(shape, False))
sigmoid_use_exp2 = bool(_RANKED_SIGMOID_EXP2.get(shape, _SIGMOID_USE_EXP2_DEFAULT))
sigmoid_use_rcp = bool(_RANKED_SIGMOID_RCP.get(shape, _USE_RCP_DEFAULT))
out_dtype = cutlass.Float16
return (
_compile_dual_gemm_kernel(
out_dtype,
mma_tiler_mn,
cluster_shape_mn,
ab_stage_cap,
occupancy,
force_tmem_512,
store_pipe,
sigmoid_use_exp2,
sigmoid_use_rcp,
),
out_dtype,
)
def _get_rank_gemm_kernel(out_dtype: Type[cutlass.Numeric], m: int):
mma_tiler_mn = (128, 64)
if m == 256:
cluster_shape_mn = (2, 2)
else:
cluster_shape_mn = (4, 1)
return _compile_gemm_kernel(out_dtype, mma_tiler_mn, cluster_shape_mn), out_dtype
def _get_safe_gemm_kernel(out_dtype: Type[cutlass.Numeric]):
mma_tiler_mn = (128, 64)
cluster_shape_mn = (1, 1)
return _compile_gemm_kernel(out_dtype, mma_tiler_mn, cluster_shape_mn), out_dtype
def custom_kernel(data):
a, b1, b2, _sfa, _sfb1, _sfb2, sfa_p, sfb1_p, sfb2_p, c = data
m, k_half, l = a.shape
n, _, _ = b1.shape
k = k_half * 2
mi = int(m)
ni = int(n)
ki = int(k)
li = int(l)
if _is_rank_shape(mi, ni, ki, li):
compiled, out_dtype = _get_rank_dual_kernel(mi, ni, ki)
_run_dual_gemm(compiled, out_dtype, a, b1, b2, sfa_p, sfb1_p, sfb2_p, c)
return c
else:
out_dtype = cutlass.Float32
g1 = _get_buf("g1", c.shape, c.stride(), c.device, torch.float32)
g2 = _get_buf("g2", c.shape, c.stride(), c.device, torch.float32)
compiled, out_dtype = _get_safe_gemm_kernel(out_dtype)
_run_gemm(compiled, out_dtype, a, b1, sfa_p, sfb1_p, g1)
_run_gemm(compiled, out_dtype, a, b2, sfa_p, sfb2_p, g2)
F.silu(g1, inplace=True)
g1.mul_(g2)
c.copy_(g1)
del g1, g2
return c
__all__ = ["custom_kernel"]
scrolls · 2790 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 375747.
⋯ 2 unchanged linesfrom typing import Optional, Tuple, Type, Union+import os- os.environ.setdefault("CUTE_DSL_ARCH", "sm_100a")- os.environ.setdefault("TARGET_SM_ARCH", "sm_100a")- os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")- _DEBUG_FORCE_EPILOGUE_SYNC = os.environ.get("NVFP4_DEBUG_FORCE_EPILOGUE_SYNC", "0") == "1"- _SIGMOID_USE_EXP2 = os.environ.get("NVFP4_SIGMOID_USE_EXP2", "0") == "1"- _USE_RCP = os.environ.get("NVFP4_USE_RCP", "0") == "1"- _FORCE_TMEM_512 = os.environ.get("NVFP4_FORCE_TMEM_512", "0") == "1"- _ENABLE_STORE_PIPE = os.environ.get("NVFP4_ENABLE_STORE_PIPE", "0") == "1"- _RANK_CLUSTER_MODE = os.environ.get("NVFP4_RANK_CLUSTER_MODE", "0")- _DIAG_MODE = os.environ.get("NVFP4_DIAG_MODE", "0").lower()- _DIAG_ACC1 = _DIAG_MODE == "acc1"- _DIAG_SINGLE = _DIAG_MODE == "single"- _DIAG_META = _DIAG_MODE == "meta"+ _TARGET_ARCH = "sm_100a"++ _SIGMOID_USE_EXP2_DEFAULT = True+ _USE_RCP_DEFAULT = True+ _SIGMOID_ENABLE_SAT = False+++ os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"+import torchimport torch.nn.functional as F⋯ 21 unchanged lines_HAS_RCP = _CUTE_RCP is not None- def _sigmoid_exp_neg(x):- if _SIGMOID_USE_EXP2 and _CUTE_EXP2 is not None:+ def _sigmoid_exp_neg(x, use_exp2: bool):+ if use_exp2 and _CUTE_EXP2 is not None:return _CUTE_EXP2(x * cutlass.Float32(_LOG2E))return cmath.exp(x, fastmath=True)+ def _silu_mul_acc(acc1, acc2, one, zero, use_exp2: bool, use_rcp: bool):+ if cutlass.const_expr(_SIGMOID_ENABLE_SAT):+ pos = acc1 > cutlass.Float32(_SIGMOID_POS_TH)+ neg = acc1 < cutlass.Float32(_SIGMOID_NEG_TH)++ den = one + _sigmoid_exp_neg(zero - acc1, use_exp2)+ if cutlass.const_expr(use_rcp and _HAS_RCP):+ inv = _CUTE_RCP(den)+ else:+ inv = cmath.rsqrt(den, fastmath=True)+ inv = inv * inv++ inv = cute.where(pos, one, inv)+ inv = cute.where(neg, zero, inv)+ return (acc1 * acc2) * inv++ den = one + _sigmoid_exp_neg(zero - acc1, use_exp2)+ if cutlass.const_expr(use_rcp and _HAS_RCP):+ inv = _CUTE_RCP(den)+ else:+ inv = cmath.rsqrt(den, fastmath=True)+ inv = inv * inv+ return (acc1 * acc2) * inv++class Sm100BlockScaledDenseGemmKernel:def __init__(self,⋯ 117 unchanged linesab_stage_cap = 0if self.mma_tiler_mn == (128, 128):ab_stage_cap = 4- env_ab_cap = int(os.environ.get("NVFP4_AB_STAGE_CAP", "0"))- if env_ab_cap > 0:- ab_stage_cap = env_ab_capif ab_stage_cap > 0 and self.num_ab_stage > ab_stage_cap:self.num_ab_stage = ab_stage_cap⋯ 812 unchanged linesmbar_helpers_bytes = 1024c_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 - (mbar_helpers_bytes + c_bytes)) // ab_bytes_per_stage- num_c_stage += (- smem_capacity - ab_bytes_per_stage * num_ab_stage - (mbar_helpers_bytes + c_bytes)- ) // (c_bytes_per_stage)+ smem_budget = smem_capacity // max(1, int(occupancy))+ budget_for_ab = smem_budget - (mbar_helpers_bytes + c_bytes)+ num_ab_stage = budget_for_ab // ab_bytes_per_stage+ if num_ab_stage < 1:+ num_ab_stage = 1+ leftover = smem_budget - ab_bytes_per_stage * num_ab_stage - (mbar_helpers_bytes + c_bytes)+ if leftover > 0:+ num_c_stage += leftover // (c_bytes_per_stage)+ if num_c_stage < 1:+ num_c_stage = 1return num_acc_stage, num_ab_stage, num_c_stage@staticmethod⋯ 17 unchanged linescluster_shape_mn: Tuple[int, int],c_dtype: Type[cutlass.Numeric],ab_stage_cap: int = 0,+ occupancy: int = 1,+ force_tmem_512: bool = False,+ store_pipe: bool = False,+ sigmoid_use_exp2: bool = True,+ sigmoid_use_rcp: bool = True,):self.ab_dtype = cutlass.Float4E2M1FNself.sf_dtype = cutlass.Float8E4M3FN⋯ 12 unchanged linesself.cluster_shape_mn = cluster_shape_mnself.ab_stage_cap = ab_stage_cap- self.occupancy = 1+ occ = int(occupancy)+ self.occupancy = 1 if occ <= 0 else occ+ self.force_tmem_512 = bool(force_tmem_512)+ self.store_pipe = bool(store_pipe)+ self.sigmoid_use_exp2 = bool(sigmoid_use_exp2)+ self.sigmoid_use_rcp = bool(sigmoid_use_rcp)self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")self.num_tmem_alloc_cols = 512⋯ 90 unchanged linesif ab_stage_cap > 0 and self.num_ab_stage > ab_stage_cap:self.num_ab_stage = ab_stage_cap- if self.cta_tile_shape_mnk[1] == 64 and self.num_c_stage > 4:- self.num_c_stage = 4-self.prefetch_stage = self.num_ab_stageself.a_smem_layout_staged = sm100_utils.make_smem_layout_a(⋯ 78 unchanged linestotal_cols = acc_cols_0 + acc_cols_1 + sfa_cols + sfb1_cols + sfb2_colsif self.mma_tiler_mn == (128, 128):total_cols = acc_cols_0 + acc_cols_1 + sfa_cols + sfb1_cols- if self.cta_tile_shape_mnk[1] == 64 and _FORCE_TMEM_512:+ if self.cta_tile_shape_mnk[1] == 64:total_cols += 2needed_cols = int(total_cols)if needed_cols <= 32:⋯ 9 unchanged lineselse:raise ValueError(f"TMEM cols budget exceeded: needed_cols={needed_cols}")- if _FORCE_TMEM_512 and self.mma_tiler_mn == (128, 128) and self.num_tmem_alloc_cols < 512:+ if self.force_tmem_512 and self.mma_tiler_mn == (128, 128) and self.num_tmem_alloc_cols < 512:self.num_tmem_alloc_cols = 512@cute.jit⋯ 795 unchanged linessubtile_cnt = cute.size(tTR_tAcc_1.shape, mode=[3])one = cutlass.Float32(1.0)zero = cutlass.Float32(0.0)- batch = 1 if _DEBUG_FORCE_EPILOGUE_SYNC else self.num_c_stage+ batch = self.num_c_stagetail = subtile_cnt % batchmain_cnt = subtile_cnt - tail⋯ 5 unchanged linesproducer_group=c_producer_group,)- if cutlass.const_expr(- _ENABLE_STORE_PIPE- or self.cta_tile_shape_mnk[1] == 128- or (self.cta_tile_shape_mnk[1] == 64 and subtile_cnt >= 4)- ):+ if cutlass.const_expr(self.store_pipe or self.cta_tile_shape_mnk[1] == 128):if warp_idx == self.epilog_warp_id[0]:c_pipeline.producer_acquire()self.epilog_sync_barrier.sync()- for base in range(0, main_cnt, batch):- for off in range(batch):- subtile_idx = base + off- tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, subtile_idx)]- cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)+ if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):+ for base in range(0, main_cnt, batch):+ for off in range(batch):+ subtile_idx = base + off+ c_buf = subtile_idx % self.num_c_stage+ if base != 0 and off == 0 and warp_idx != self.epilog_warp_id[0]:+ cute.copy(+ tiled_copy_r2s,+ tRS_rC,+ tRS_sC[(None, None, None, c_buf)],+ )+ else:+ tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, subtile_idx)]+ cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)- acc1 = tTR_rAcc_1.load()- if cutlass.const_expr(_DIAG_ACC1):- out = acc1- else:+ acc1 = tTR_rAcc_1.load()+ tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, subtile_idx)]+ cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)+ acc2 = tTR_rAcc_2.load()+ out = _silu_mul_acc(+ acc1,+ acc2,+ one,+ zero,+ self.sigmoid_use_exp2,+ self.sigmoid_use_rcp,+ )+ tRS_rC.store(out.to(self.c_dtype))+ cute.copy(+ tiled_copy_r2s,+ tRS_rC,+ tRS_sC[(None, None, None, c_buf)],+ )++ cute.arch.fence_view_async_shared()+ self.epilog_sync_barrier.sync()+ if base + batch < main_cnt and warp_idx != self.epilog_warp_id[0]:+ subtile_idx = base + batch+ tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, subtile_idx)]+ cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)++ acc1 = tTR_rAcc_1.load()tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, subtile_idx)]cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)acc2 = tTR_rAcc_2.load()- if cutlass.const_expr(- self.cta_tile_shape_mnk[1] == 64 and _CUTE_EXP2 is not None- ):- den = one + _CUTE_EXP2(- (zero - acc1) * cutlass.Float32(_LOG2E)+ out = _silu_mul_acc(+ acc1,+ acc2,+ one,+ zero,+ self.sigmoid_use_exp2,+ self.sigmoid_use_rcp,+ )+ tRS_rC.store(out.to(self.c_dtype))+ if warp_idx == self.epilog_warp_id[0]:+ for off in range(batch):+ subtile_idx = base + off+ c_buf = subtile_idx % self.num_c_stage+ cute.copy(+ tma_atom_c,+ bSG_sC[(None, c_buf)],+ bSG_gC[(None, subtile_idx)],)- else:- den = one + _sigmoid_exp_neg(zero - acc1)- if cutlass.const_expr(- (self.cta_tile_shape_mnk[1] == 64 and _HAS_RCP)- or (_USE_RCP and _HAS_RCP)- ):- inv = _CUTE_RCP(den)- else:- inv = cmath.rsqrt(den, fastmath=True)- inv = inv * inv- out = (acc1 * acc2) * inv- tRS_rC.store(out.to(self.c_dtype))- cute.copy(- tiled_copy_r2s,- tRS_rC,- tRS_sC[(None, None, None, subtile_idx)],- )+ c_pipeline.producer_commit()+ c_pipeline.producer_acquire()+ self.epilog_sync_barrier.sync()- cute.arch.fence_view_async_shared()- self.epilog_sync_barrier.sync()- if warp_idx == self.epilog_warp_id[0]:- for off in range(batch):- store_idx = base + off+ if tail != 0:+ base = main_cnt+ for off in range(tail):+ subtile_idx = base + off+ c_buf = subtile_idx % self.num_c_stage+ tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, subtile_idx)]+ cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)++ acc1 = tTR_rAcc_1.load()+ tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, subtile_idx)]+ cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)+ acc2 = tTR_rAcc_2.load()+ out = _silu_mul_acc(+ acc1,+ acc2,+ one,+ zero,+ self.sigmoid_use_exp2,+ self.sigmoid_use_rcp,+ )+ tRS_rC.store(out.to(self.c_dtype))cute.copy(- tma_atom_c,- bSG_sC[(None, store_idx)],- bSG_gC[(None, store_idx)],+ tiled_copy_r2s,+ tRS_rC,+ tRS_sC[(None, None, None, c_buf)],)- c_pipeline.producer_commit()- c_pipeline.producer_acquire()- self.epilog_sync_barrier.sync()- if tail != 0:- base = main_cnt- for off in range(tail):- subtile_idx = base + off- tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, subtile_idx)]- cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)+ cute.arch.fence_view_async_shared()+ self.epilog_sync_barrier.sync()+ if warp_idx == self.epilog_warp_id[0]:+ for off in range(tail):+ subtile_idx = base + off+ c_buf = subtile_idx % self.num_c_stage+ cute.copy(+ tma_atom_c,+ bSG_sC[(None, c_buf)],+ bSG_gC[(None, subtile_idx)],+ )+ c_pipeline.producer_commit()+ c_pipeline.producer_acquire()+ self.epilog_sync_barrier.sync()- acc1 = tTR_rAcc_1.load()- if cutlass.const_expr(_DIAG_ACC1):- out = acc1- else:+ self.epilog_sync_barrier.sync()+ if warp_idx == self.epilog_warp_id[0]:+ c_pipeline.producer_tail()+ self.epilog_sync_barrier.sync()+ else:+ for base in range(0, main_cnt, batch):+ for off in range(batch):+ subtile_idx = base + off+ tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, subtile_idx)]+ cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)++ acc1 = tTR_rAcc_1.load()tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, subtile_idx)]cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)acc2 = tTR_rAcc_2.load()- if cutlass.const_expr(- self.cta_tile_shape_mnk[1] == 64 and _CUTE_EXP2 is not None- ):- den = one + _CUTE_EXP2(- (zero - acc1) * cutlass.Float32(_LOG2E)+ out = _silu_mul_acc(+ acc1,+ acc2,+ one,+ zero,+ self.sigmoid_use_exp2,+ self.sigmoid_use_rcp,+ )+ tRS_rC.store(out.to(self.c_dtype))+ cute.copy(+ tiled_copy_r2s,+ tRS_rC,+ tRS_sC[(None, None, None, subtile_idx)],+ )++ cute.arch.fence_view_async_shared()+ self.epilog_sync_barrier.sync()+ if warp_idx == self.epilog_warp_id[0]:+ for off in range(batch):+ store_idx = base + off+ cute.copy(+ tma_atom_c,+ bSG_sC[(None, store_idx)],+ bSG_gC[(None, store_idx)],)- else:- den = one + _sigmoid_exp_neg(zero - acc1)- if cutlass.const_expr(- (self.cta_tile_shape_mnk[1] == 64 and _HAS_RCP)- or (_USE_RCP and _HAS_RCP)- ):- inv = _CUTE_RCP(den)- else:- inv = cmath.rsqrt(den, fastmath=True)- inv = inv * inv- out = (acc1 * acc2) * inv- tRS_rC.store(out.to(self.c_dtype))- cute.copy(- tiled_copy_r2s,- tRS_rC,- tRS_sC[(None, None, None, subtile_idx)],- )+ c_pipeline.producer_commit()+ c_pipeline.producer_acquire()+ self.epilog_sync_barrier.sync()- cute.arch.fence_view_async_shared()- self.epilog_sync_barrier.sync()- if warp_idx == self.epilog_warp_id[0]:+ if tail != 0:+ base = main_cntfor off in range(tail):- store_idx = base + off+ subtile_idx = base + off+ tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, subtile_idx)]+ cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)++ acc1 = tTR_rAcc_1.load()+ tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, subtile_idx)]+ cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)+ acc2 = tTR_rAcc_2.load()+ out = _silu_mul_acc(+ acc1,+ acc2,+ one,+ zero,+ self.sigmoid_use_exp2,+ self.sigmoid_use_rcp,+ )+ tRS_rC.store(out.to(self.c_dtype))cute.copy(- tma_atom_c,- bSG_sC[(None, store_idx)],- bSG_gC[(None, store_idx)],+ tiled_copy_r2s,+ tRS_rC,+ tRS_sC[(None, None, None, subtile_idx)],)- c_pipeline.producer_commit()- c_pipeline.producer_acquire()- self.epilog_sync_barrier.sync()- self.epilog_sync_barrier.sync()- if warp_idx == self.epilog_warp_id[0]:- c_pipeline.producer_tail()- self.epilog_sync_barrier.sync()- else:- for subtile_idx in range(subtile_cnt):- tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, subtile_idx)]- cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)-- acc1 = tTR_rAcc_1.load()- if cutlass.const_expr(_DIAG_ACC1):- out = acc1- else:- tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, subtile_idx)]- cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)- acc2 = tTR_rAcc_2.load()- if cutlass.const_expr(- self.cta_tile_shape_mnk[1] == 64 and _CUTE_EXP2 is not None- ):- den = one + _CUTE_EXP2((zero - acc1) * cutlass.Float32(_LOG2E))- else:- den = one + _sigmoid_exp_neg(zero - acc1)- if cutlass.const_expr(- (self.cta_tile_shape_mnk[1] == 64 and _HAS_RCP)- or (_USE_RCP and _HAS_RCP)- ):- inv = _CUTE_RCP(den)- else:- inv = cmath.rsqrt(den, fastmath=True)- inv = inv * inv- out = (acc1 * acc2) * inv- tRS_rC.store(out.to(self.c_dtype))- cute.copy(- tiled_copy_r2s,- tRS_rC,- tRS_sC[(None, None, None, subtile_idx)],- )- if (subtile_idx + 1) % batch == 0:cute.arch.fence_view_async_shared()self.epilog_sync_barrier.sync()if warp_idx == self.epilog_warp_id[0]:- for off in range(batch):- store_idx = subtile_idx - (batch - 1 - off)+ for off in range(tail):+ store_idx = base + offcute.copy(tma_atom_c,bSG_sC[(None, store_idx)],bSG_gC[(None, store_idx)],)+ c_pipeline.producer_commit()+ c_pipeline.producer_acquire()+ self.epilog_sync_barrier.sync()+ self.epilog_sync_barrier.sync()+ if warp_idx == self.epilog_warp_id[0]:+ c_pipeline.producer_tail()+ self.epilog_sync_barrier.sync()+ else:+ tTR_rAcc_1_b = cute.make_fragment(tTR_rAcc_1.shape, self.acc_dtype)+ tTR_rAcc_2_b = cute.make_fragment(tTR_rAcc_2.shape, self.acc_dtype)++ if subtile_cnt > 0:+ tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, 0)]+ cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)+ tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, 0)]+ cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)++ for subtile_idx in range(subtile_cnt):+ cur_is_0 = (subtile_idx & 1) == 0+ next_idx = subtile_idx + 1++ if next_idx < subtile_cnt:+ if cur_is_0:+ tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, next_idx)]+ cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1_b)+ tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, next_idx)]+ cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2_b)+ else:+ tTR_tAcc_mn_1 = tTR_tAcc_1[(None, None, None, next_idx)]+ cute.copy(tiled_copy_t2r_1, tTR_tAcc_mn_1, tTR_rAcc_1)+ tTR_tAcc_mn_2 = tTR_tAcc_2[(None, None, None, next_idx)]+ cute.copy(tiled_copy_t2r_2, tTR_tAcc_mn_2, tTR_rAcc_2)++ if cur_is_0:+ acc1 = tTR_rAcc_1.load()+ acc2 = tTR_rAcc_2.load()+ out = _silu_mul_acc(+ acc1,+ acc2,+ one,+ zero,+ self.sigmoid_use_exp2,+ self.sigmoid_use_rcp,+ )+ tRS_rC.store(out.to(self.c_dtype))+ else:+ acc1 = tTR_rAcc_1_b.load()+ acc2 = tTR_rAcc_2_b.load()+ out = _silu_mul_acc(+ acc1,+ acc2,+ one,+ zero,+ self.sigmoid_use_exp2,+ self.sigmoid_use_rcp,+ )+ tRS_rC.store(out.to(self.c_dtype))+ cute.copy(+ tiled_copy_r2s,+ tRS_rC,+ tRS_sC[(None, None, None, subtile_idx)],+ )++ if (subtile_idx + 1) % batch == 0:+ cute.arch.fence_view_async_shared()+ self.epilog_sync_barrier.sync()+ if warp_idx == self.epilog_warp_id[0]:+ for off in range(batch):+ store_idx = subtile_idx - (batch - 1 - off)+ cute.copy(+ tma_atom_c,+ bSG_sC[(None, store_idx)],+ bSG_gC[(None, store_idx)],+ )+if tail != 0:cute.arch.fence_view_async_shared()self.epilog_sync_barrier.sync()⋯ 66 unchanged linesmbar_helpers_bytes = 1024c_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 - (mbar_helpers_bytes + c_bytes)) // ab_bytes_per_stage- num_c_stage += (- smem_capacity - ab_bytes_per_stage * num_ab_stage - (mbar_helpers_bytes + c_bytes)- ) // (c_bytes_per_stage)+ smem_budget = smem_capacity // max(1, int(occupancy))+ budget_for_ab = smem_budget - (mbar_helpers_bytes + c_bytes)+ num_ab_stage = budget_for_ab // ab_bytes_per_stage+ if num_ab_stage < 1:+ num_ab_stage = 1+ leftover = smem_budget - ab_bytes_per_stage * num_ab_stage - (mbar_helpers_bytes + c_bytes)+ if leftover > 0:+ num_c_stage += leftover // (c_bytes_per_stage)+ if num_c_stage < 1:+ num_c_stage = 1return num_acc_stage, num_ab_stage, num_c_stage_GEMM_KERNEL_CACHE = {}_DUAL_GEMM_KERNEL_CACHE = {}+ _CUDA_CACHE_CLEARED = False+ _BUF_CACHE = {}+++ def _maybe_empty_cuda_cache_once() -> None:+ global _CUDA_CACHE_CLEARED+ if not _CUDA_CACHE_CLEARED:+ torch.cuda.empty_cache()+ _CUDA_CACHE_CLEARED = True++def _get_buf(tag: str, shape, stride, device, dtype: torch.dtype) -> torch.Tensor:- return torch.empty_strided(shape, stride, device=device, dtype=dtype)+++ if isinstance(shape, torch.Size):+ shape = tuple(int(x) for x in shape)+ else:+ shape = tuple(int(x) for x in shape)+ if isinstance(stride, torch.Size):+ stride = tuple(int(x) for x in stride)+ else:+ stride = tuple(int(x) for x in stride)++ if (+ len(shape) == 3+ and shape[2] == 1+ and len(stride) == 3+ and stride[1] == 1+ and stride[2] == 1+ and stride[0] >= shape[1]+ ):+ m, n, l = shape+ key = (tag, int(stride[0]), int(n), device, dtype)+ buf = _BUF_CACHE.get(key)+ if buf is None or buf.shape[0] < m:++ buf = torch.empty_strided((m, n, l), stride, device=device, dtype=dtype)+ _BUF_CACHE[key] = buf+ return buf[:m, :n, :l]++ key = (tag, shape, stride, device, dtype)+ buf = _BUF_CACHE.get(key)+ if buf is None:+ buf = torch.empty_strided(shape, stride, device=device, dtype=dtype)+ _BUF_CACHE[key] = buf+ return buf++def _compile_gemm_kernel(out_dtype: Type[cutlass.Numeric], mma_tiler_mn, cluster_shape_mn):key = (out_dtype, mma_tiler_mn, cluster_shape_mn)cached = _GEMM_KERNEL_CACHE.get(key)if cached is not None:return cached- target_arch = os.environ.get("CUTE_DSL_ARCH", "sm_100a")+ target_arch = _TARGET_ARCHcutlass.cuda.initialize_cuda_context()a_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)⋯ 7 unchanged lineskernel, a_ptr, b_ptr, sfa_ptr, sfb_ptr, c_ptr, 0, 0, 0, 0)_GEMM_KERNEL_CACHE[key] = compiled+ _maybe_empty_cuda_cache_once()return compiled⋯ 2 unchanged linesmma_tiler_mn,cluster_shape_mn,ab_stage_cap: int,+ occupancy: int,+ force_tmem_512: bool,+ store_pipe: bool,+ sigmoid_use_exp2: bool,+ sigmoid_use_rcp: bool,):- key = (out_dtype, mma_tiler_mn, cluster_shape_mn, int(ab_stage_cap))+ key = (+ out_dtype,+ mma_tiler_mn,+ cluster_shape_mn,+ int(ab_stage_cap),+ int(occupancy),+ bool(force_tmem_512),+ bool(store_pipe),+ bool(sigmoid_use_exp2),+ bool(sigmoid_use_rcp),+ bool(_SIGMOID_ENABLE_SAT),+ )cached = _DUAL_GEMM_KERNEL_CACHE.get(key)if cached is not None:return cached- target_arch = os.environ.get("CUTE_DSL_ARCH", "sm_100a")+ target_arch = _TARGET_ARCHcutlass.cuda.initialize_cuda_context()a_ptr = make_ptr(ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16)⋯ 9 unchanged linescluster_shape_mn,c_dtype=out_dtype,ab_stage_cap=int(ab_stage_cap),+ occupancy=int(occupancy),+ force_tmem_512=bool(force_tmem_512),+ store_pipe=bool(store_pipe),+ sigmoid_use_exp2=bool(sigmoid_use_exp2),+ sigmoid_use_rcp=bool(sigmoid_use_rcp),)compiled = cute.compile[cute.GPUArch(target_arch)](kernel,⋯ 10 unchanged lines0,)_DUAL_GEMM_KERNEL_CACHE[key] = compiled- if _DIAG_META:- raise RuntimeError(- f"dual_meta mma_tiler_mn={mma_tiler_mn} cluster_shape_mn={cluster_shape_mn} "- f"ab_stage_cap={int(ab_stage_cap)} num_ab_stage={getattr(kernel, 'num_ab_stage', None)} "- f"num_acc_stage={getattr(kernel, 'num_acc_stage', None)} num_c_stage={getattr(kernel, 'num_c_stage', None)} "- f"num_tmem_alloc_cols={getattr(kernel, 'num_tmem_alloc_cols', None)} epi_tile={getattr(kernel, 'epi_tile', None)} "- f"cta_tile_shape_mnk={getattr(kernel, 'cta_tile_shape_mnk', None)}"- )+ _maybe_empty_cuda_cache_once()return compiled⋯ 55 unchanged lines}_RANKED_CLUSTER_SHAPE_V1 = {- (256, 4096, 7168): (1, 1),- (512, 4096, 7168): (2, 2),- (256, 3072, 4096): (2, 2),- (512, 3072, 7168): (2, 2),- }-- _RANKED_CLUSTER_SHAPE_V2 = {- (256, 4096, 7168): (2, 1),- (512, 4096, 7168): (2, 2),- (256, 3072, 4096): (2, 1),- (512, 3072, 7168): (2, 2),- }-- _RANKED_CLUSTER_SHAPE_V3 = {- (256, 4096, 7168): (1, 2),- (512, 4096, 7168): (2, 2),- (256, 3072, 4096): (1, 2),- (512, 3072, 7168): (2, 2),- }-- _RANKED_CLUSTER_SHAPE_V4 = {(256, 4096, 7168): (2, 2),(512, 4096, 7168): (4, 1),(256, 3072, 4096): (2, 2),(512, 3072, 7168): (4, 1),}- _RANKED_CLUSTER_SHAPE = _RANKED_CLUSTER_SHAPE_V0- if _RANK_CLUSTER_MODE == "1":- _RANKED_CLUSTER_SHAPE = _RANKED_CLUSTER_SHAPE_V1- elif _RANK_CLUSTER_MODE == "2":- _RANKED_CLUSTER_SHAPE = _RANKED_CLUSTER_SHAPE_V2- elif _RANK_CLUSTER_MODE == "3":- _RANKED_CLUSTER_SHAPE = _RANKED_CLUSTER_SHAPE_V3- elif _RANK_CLUSTER_MODE == "4":- _RANKED_CLUSTER_SHAPE = _RANKED_CLUSTER_SHAPE_V4+ _RANKED_DUAL_VARIANT = {+ (256, 4096, 7168): 0,+ (512, 4096, 7168): 1,+ (256, 3072, 4096): 0,+ (512, 3072, 7168): 1,+ }_RANKED_MMA_TILER_MN = {- (256, 4096, 7168): (128, 128),+ (256, 4096, 7168): (128, 64),(512, 4096, 7168): (128, 128),(256, 3072, 4096): (128, 64),(512, 3072, 7168): (128, 128),}_RANKED_AB_STAGE_CAP = {- (256, 4096, 7168): 0,- (512, 4096, 7168): 0,- (256, 3072, 4096): 4,- (512, 3072, 7168): 0,+ (256, 4096, 7168): 99,+ (512, 4096, 7168): 99,+ (256, 3072, 4096): 99,+ (512, 3072, 7168): 99,}+ _RANKED_OCCUPANCY = {+ (256, 4096, 7168): 1,+ (512, 4096, 7168): 1,+ (256, 3072, 4096): 1,+ (512, 3072, 7168): 1,+ }+ _RANKED_STORE_PIPE = {+ (256, 4096, 7168): True,+ (512, 4096, 7168): False,+ (256, 3072, 4096): False,+ (512, 3072, 7168): False,+ }++ _RANKED_FORCE_TMEM_512 = {+ (256, 4096, 7168): False,+ (512, 4096, 7168): False,+ (256, 3072, 4096): False,+ (512, 3072, 7168): True,+ }++ _RANKED_SIGMOID_EXP2 = {+ (256, 4096, 7168): _SIGMOID_USE_EXP2_DEFAULT,+ (512, 4096, 7168): _SIGMOID_USE_EXP2_DEFAULT,+ (256, 3072, 4096): _SIGMOID_USE_EXP2_DEFAULT,+ (512, 3072, 7168): False,+ }++ _RANKED_SIGMOID_RCP = {+ (256, 4096, 7168): _USE_RCP_DEFAULT,+ (512, 4096, 7168): _USE_RCP_DEFAULT,+ (256, 3072, 4096): _USE_RCP_DEFAULT,+ (512, 3072, 7168): True,+ }++def _get_rank_dual_kernel(m: int, n: int, k: int):- mma_tiler_mn = _RANKED_MMA_TILER_MN.get((m, n, k), (128, 64))- cluster_shape_mn = _RANKED_CLUSTER_SHAPE.get((m, n, k))- if cluster_shape_mn is None:- cluster_shape_mn = (2, 2)- ab_stage_cap = int(_RANKED_AB_STAGE_CAP.get((m, n, k), 0))+ shape = (m, n, k)+ variant = int(_RANKED_DUAL_VARIANT.get(shape, 0))+ mma_tiler_mn = _RANKED_MMA_TILER_MN.get(shape, (128, 64))+ cluster_shape_mn = (_RANKED_CLUSTER_SHAPE_V1 if variant == 1 else _RANKED_CLUSTER_SHAPE_V0).get(shape, (2, 2))+ ab_stage_cap = int(_RANKED_AB_STAGE_CAP.get(shape, 0))+ occupancy = int(_RANKED_OCCUPANCY.get(shape, 1))+ store_pipe = bool(_RANKED_STORE_PIPE.get(shape, False))+ force_tmem_512 = bool(_RANKED_FORCE_TMEM_512.get(shape, False))+ sigmoid_use_exp2 = bool(_RANKED_SIGMOID_EXP2.get(shape, _SIGMOID_USE_EXP2_DEFAULT))+ sigmoid_use_rcp = bool(_RANKED_SIGMOID_RCP.get(shape, _USE_RCP_DEFAULT))out_dtype = cutlass.Float16return (- _compile_dual_gemm_kernel(out_dtype, mma_tiler_mn, cluster_shape_mn, ab_stage_cap),+ _compile_dual_gemm_kernel(+ out_dtype,+ mma_tiler_mn,+ cluster_shape_mn,+ ab_stage_cap,+ occupancy,+ force_tmem_512,+ store_pipe,+ sigmoid_use_exp2,+ sigmoid_use_rcp,+ ),out_dtype,)⋯ 26 unchanged linesli = int(l)if _is_rank_shape(mi, ni, ki, li):- if _DIAG_SINGLE:- compiled, out_dtype = _get_rank_gemm_kernel(cutlass.Float16, mi)- _run_gemm(compiled, out_dtype, a, b1, sfa_p, sfb1_p, c)- return ccompiled, out_dtype = _get_rank_dual_kernel(mi, ni, ki)_run_dual_gemm(compiled, out_dtype, a, b1, b2, sfa_p, sfb1_p, sfb2_p, c)return c
scrolls · 876 diff lines total
Best evidence level for this revision: reported
JSON