submission 275665
guaguabear · python · License unknown
Use it
Vendorable · source mirrored · license unknownView source →
No package. Vendor the mirrored source: 2352 lines, June 9 Researcher Reciprocity License v1.0.
test.py
curl "https://kernelindex.com/api/v1/implementations/kernelbot-nvfp4-dual-gemm-275665?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:d52ccb373353ec2d787ce77c41d43ebd7a3069f04e2e901438ef1ca02a415d63
license declaredunknown
license concludedunknown
authorsguaguabear
imported2026-08-15
Techniques
Extracted from the mirrored source by pattern, never inferred. Each row cites its line.
fp4
PyTorch reference implementation of NVFP4 block-scaled dual GEMM with silu activation,fused-epilogue
- Computing epilogue subtilembarrier
self.epilog_sync_barrier = pipeline.NamedBarrier(persistent-kernel
and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization.shared-memory
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")tcgen05
tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONEwarp-specialization
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)Kernel source
test.py2352 lines
from torch._higher_order_ops.torchbind import call_torchbind_fake
import cuda.bindings.driver as cuda
import subprocess
import sys
from task import input_t, output_t
import argparse
from typing import Type, Tuple, Union
import cuda.bindings.driver as cuda
import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.nvgpu import cpasync, tcgen05
import cutlass.torch as cutlass_torch
import cutlass.utils as utils
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils.blockscaled_layout as blockscaled_utils
from cutlass.cute.runtime import from_dlpack
from cutlass.cute.runtime import make_ptr
from cutlass.cutlass_dsl import CuTeDSL, dsl_user_op, T
from cutlass._mlir import ir
from cutlass._mlir.dialects import builtin, arith, llvm, vector
from cutlass.cute.typing import (
Int4,
Int8,
Int16,
Int32,
Float16,
Float32,
BFloat16,
Float32,
)
# 14.5
# 16.9
# 10.3
# 16.5
# problem size: tile_mn | cluster_mn | pf_dist
config_map = {
(256,4096,7168) : ((256, 64), (2, 1), 0, True),
(512,4096,7168) : ((256, 128), (2, 1), 0, False),
(256,3072,4096) : ((256, 64), (2, 1), None, True),
(512,3072,7168) : ((256, 128), (2, 1), None, True),
}
ab_dtype = cutlass.Float4E2M1FN
sf_dtype = cutlass.Float8E4M3FN
c_dtype = cutlass.Float16
sf_vec_size = 16
max_active_clusters = 148
@dsl_user_op
def silu_precise_8(src_A, src_B, appro=False, loc=None, ip=None):
inputs = []
# 提取 A0-A7
for i in range(8):
inputs.append(llvm.extractelement(src_A, arith.constant(Int32.mlir_type, i, loc=loc, ip=ip), loc=loc, ip=ip))
# 提取 B0-B7
for i in range(8):
inputs.append(llvm.extractelement(src_B, arith.constant (Int32.mlir_type, i, loc=loc, ip=ip), loc=loc, ip=ip))
# A: $8-$15, B: $16-$23, Out: $0-$7
asm = r"""
mul.f32 $0, $8, 0fBFB8AA3B; mul.f32 $1, $9, 0fBFB8AA3B; mul.f32 $2, $10, 0fBFB8AA3B; mul.f32 $3, $11, 0fBFB8AA3B;
mul.f32 $4, $12, 0fBFB8AA3B; mul.f32 $5, $13, 0fBFB8AA3B; mul.f32 $6, $14, 0fBFB8AA3B; mul.f32 $7, $15, 0fBFB8AA3B;
ex2.approx.f32 $0, $0; ex2.approx.f32 $1, $1; ex2.approx.f32 $2, $2; ex2.approx.f32 $3, $3;
ex2.approx.f32 $4, $4; ex2.approx.f32 $5, $5; ex2.approx.f32 $6, $6; ex2.approx.f32 $7, $7;
add.f32 $0, $0, 1.0; add.f32 $1, $1, 1.0; add.f32 $2, $2, 1.0; add.f32 $3, $3, 1.0;
add.f32 $4, $4, 1.0; add.f32 $5, $5, 1.0; add.f32 $6, $6, 1.0; add.f32 $7, $7, 1.0;
rcp.approx.f32 $0, $0; rcp.approx.f32 $1, $1; rcp.approx.f32 $2, $2; rcp.approx.f32 $3, $3;
rcp.approx.f32 $4, $4; rcp.approx.f32 $5, $5; rcp.approx.f32 $6, $6; rcp.approx.f32 $7, $7;
mul.f32 $0, $8, $0; mul.f32 $1, $9, $1; mul.f32 $2, $10, $2; mul.f32 $3, $11, $3;
mul.f32 $4, $12, $4; mul.f32 $5, $13, $5; mul.f32 $6, $14, $6; mul.f32 $7, $15, $7;
mul.f32 $0, $16, $0; mul.f32 $1, $17, $1; mul.f32 $2, $18, $2; mul.f32 $3, $19, $3;
mul.f32 $4, $20, $4; mul.f32 $5, $21, $5; mul.f32 $6, $22, $6; mul.f32 $7, $23, $7;
"""
if appro:
asm = r"""
mul.f32 $0, $8, 0.5; mul.f32 $1, $9, 0.5; mul.f32 $2, $10, 0.5; mul.f32 $3, $11, 0.5;
mul.f32 $4, $12, 0.5; mul.f32 $5, $13, 0.5; mul.f32 $6, $14, 0.5; mul.f32 $7, $15, 0.5;
tanh.approx.f32 $0, $0; tanh.approx.f32 $1, $1; tanh.approx.f32 $2, $2; tanh.approx.f32 $3, $3;
tanh.approx.f32 $4, $4; tanh.approx.f32 $5, $5; tanh.approx.f32 $6, $6; tanh.approx.f32 $7, $7;
add.f32 $0, $0, 1.0; add.f32 $1, $1, 1.0; add.f32 $2, $2, 1.0; add.f32 $3, $3, 1.0;
add.f32 $4, $4, 1.0; add.f32 $5, $5, 1.0; add.f32 $6, $6, 1.0; add.f32 $7, $7, 1.0;
mul.f32 $0, $0, $8; mul.f32 $1, $1, $9; mul.f32 $2, $2, $10; mul.f32 $3, $3, $11;
mul.f32 $4, $4, $12; mul.f32 $5, $5, $13; mul.f32 $6, $6, $14; mul.f32 $7, $7, $15;
mul.f32 $0, $0, 0.5; mul.f32 $1, $1, 0.5; mul.f32 $2, $2, 0.5; mul.f32 $3, $3, 0.5;
mul.f32 $4, $4, 0.5; mul.f32 $5, $5, 0.5; mul.f32 $6, $6, 0.5; mul.f32 $7, $7, 0.5;
mul.f32 $0, $16, $0; mul.f32 $1, $17, $1; mul.f32 $2, $18, $2; mul.f32 $3, $19, $3;
mul.f32 $4, $20, $4; mul.f32 $5, $21, $5; mul.f32 $6, $22, $6; mul.f32 $7, $23, $7;
"""
# asm = r"""
# mul.f32 $0, $8, 0.5; mul.f32 $1, $9, 0.5; mul.f32 $2, $10, 0.5; mul.f32 $3, $11, 0.5;
# mul.f32 $4, $12, 0.5; mul.f32 $5, $13, 0.5; mul.f32 $6, $14, 0.5; mul.f32 $7, $15, 0.5;
# {
# .reg .f32 %t<8>;
# tanh.approx.f32 %t0, $0; tanh.approx.f32 %t1, $1; tanh.approx.f32 %t2, $2; tanh.approx.f32 %t3, $3;
# tanh.approx.f32 %t4, $4; tanh.approx.f32 %t5, $5; tanh.approx.f32 %t6, $6; tanh.approx.f32 %t7, $7;
# fma.rn.f32 $0, %t0, $0, $0; fma.rn.f32 $1, %t1, $1, $1; fma.rn.f32 $2, %t2, $2, $2; fma.rn.f32 $3, %t3, $3, $3;
# fma.rn.f32 $4, %t4, $4, $4; fma.rn.f32 $5, %t5, $5, $5; fma.rn.f32 $6, %t6, $6, $6; fma.rn.f32 $7, %t7, $7, $7;
# }
# mul.f32 $0, $0, $16; mul.f32 $1, $1, $17; mul.f32 $2, $2, $18; mul.f32 $3, $3, $19;
# mul.f32 $4, $4, $20; mul.f32 $5, $5, $21; mul.f32 $6, $6, $22; mul.f32 $7, $7, $23;
# """
cons = "=&f,=&f,=&f,=&f,=&f,=&f,=&f,=&f,f,f,f,f,f,f,f,f,f,f,f,f,f,f,f,f"
# cons = "=f,=f,=f,=f,=f,=f,=f,=f,f,f,f,f,f,f,f,f,f,f,f,f,f,f,f,f"
res = llvm.inline_asm(llvm.StructType.get_literal([Float32.mlir_type]*8), inputs, asm, cons, loc=loc, ip=ip)
out = []
for i in range(8):
out.append(llvm.extractvalue(Float32.mlir_type, res, [i], loc=loc, ip=ip))
return vector.from_elements(ir.VectorType.get([8], Float32.mlir_type, loc=loc), out, loc=loc, ip=ip)
@dsl_user_op
def silu_intrinsic(vec_A, vec_B, length, appro=False, loc=None, ip=None):
src_pos = 0
vec_f32x8_type = ir.VectorType.get([8], Float32.mlir_type, loc=loc)
vec_dst_type = ir.VectorType.get([length], Float32.mlir_type, loc=loc)
vec_dst = llvm.mlir_zero(vec_dst_type, loc=loc, ip=ip)
for _ in range(length//8):
vec_f32x8_A = vector.extract_strided_slice(
vec_f32x8_type, vec_A, [src_pos], [8], [1], loc=loc, ip=ip
)
vec_f32x8_B = vector.extract_strided_slice(
vec_f32x8_type, vec_B, [src_pos], [8], [1], loc=loc, ip=ip
)
vec_dst = vector.insert_strided_slice(
silu_precise_8(vec_f32x8_A, vec_f32x8_B, appro),
vec_dst,
[src_pos],
[1],
loc=loc,
ip=ip,
)
src_pos += 8
return vec_dst
def ceil_div(a, b):
return (a + b - 1) // b
def to_blocked(input_matrix):
rows, cols = input_matrix.shape
# Please ensure rows and cols are multiples of 128 and 4 respectively
n_row_blocks = ceil_div(rows, 128)
n_col_blocks = ceil_div(cols, 4)
padded = input_matrix
blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
return rearranged.flatten()
def ref_kernel(
data: input_t,
) -> output_t:
"""
PyTorch reference implementation of NVFP4 block-scaled dual GEMM with silu activation,
C = silu(A @ B1) * (A @ B2).
"""
a_ref, b1_ref, b2_ref, sfa_ref_cpu, sfb1_ref_cpu, sfb2_ref_cpu, _, _, _, c_ref = data
# Get dimensions from MxNxL layout
m, n, l = c_ref.shape
# Call torch._scaled_mm to compute the GEMV result
ref1 = torch.empty(
(l, m, n),
dtype=torch.float32,
device="cuda",
).permute(1, 2, 0)
ref2 = torch.empty(
(l, m, n),
dtype=torch.float32,
device="cuda",
).permute(1, 2, 0)
for l_idx in range(l):
# Convert the scale factor tensor to blocked format
scale_a = to_blocked(sfa_ref_cpu[:, :, l_idx])
scale_b1 = to_blocked(sfb1_ref_cpu[:, :, l_idx])
scale_b2 = to_blocked(sfb2_ref_cpu[:, :, l_idx])
# (m, k) @ (n, k).T -> (m, n)
res1 = torch._scaled_mm(
a_ref[:, :, l_idx],
b1_ref[:, :, l_idx].transpose(0, 1),
scale_a.cuda(),
scale_b1.cuda(),
bias=None,
out_dtype=torch.float32,
)
ref1[:, :, l_idx] = res1
res2 = torch._scaled_mm(
a_ref[:, :, l_idx],
b2_ref[:, :, l_idx].transpose(0, 1),
scale_a.cuda(),
scale_b2.cuda(),
bias=None,
out_dtype=torch.float32,
)
ref2[:, :, l_idx] = res2
# Do silu on the first GEMM result and multiply with the second GEMM result
c_ref = (torch.nn.functional.silu(ref1) * ref2).to(torch.float16)
return c_ref
class Sm100BlockScaledPersistentDualGemmKernel:
"""This class implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types
and architectural features specific to Blackwell GPUs with persistent tile scheduling and warp specialization.
:param sf_vec_size: Scalefactor vector size.
:type sf_vec_size: int
:param mma_tiler_mn: Shape of the Matrix Multiply-Accumulate (MMA) tile (M,N)
:type mma_tiler_mn: Tuple[int, int]
:param cluster_shape_mn: Cluster dimensions (M,N) for parallel processing
:type cluster_shape_mn: Tuple[int, int]
:note: In current version, A and B tensor must have the same data type
- i.e., Float8E4M3FN for A and Float8E5M2 for B is not supported
:note: Supported combinations of A/B data types, SF data typs and SF vector size:
- MXF8: A/B: Float8E5M2/Float8E4M3FN + SF: Float8E8M0FNU + sf_vec_size: 32
- MXF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU + sf_vec_size: 32
- NVF4: A/B: Float4E2M1FN + SF: Float8E8M0FNU/Float8E4M3FN + sf_vec_size: 16
:note: Supported accumulator data types:
- Float32
:note: Supported C data types:
- Float32
- Float16/BFloat16
- Float8E4M3FN/Float8E5M2
:note: Constraints:
- MMA tiler M must be 128 or 256 (use_2cta_instrs)
- MMA tiler N must be 64/128/192/256
- Cluster shape M must be multiple of 2 if Mma tiler M is 256
- Cluster shape M/N must be positive and power of 2, total cluster size <= 16
- Also, Cluster shape M/N must be <= 4 for scale factor multicasts due to limited size of scale factors
Example:
>>> gemm = Sm100BlockScaledPersistentDenseGemmKernel(
... sf_vec_size=16,
... mma_tiler_mn=(256, 128),
... cluster_shape_mn=(2, 1)
... )
"""
def __init__(
self,
sf_vec_size: int,
problem_size: Tuple[int, int, int],
):
"""Initializes the configuration for a Blackwell dense GEMM kernel with TMA prefetch support.
This configuration includes several key aspects:
1. MMA Instruction Settings (tcgen05):
- acc_dtype: Data types for MMA accumulator, always set to Float32
- sf_vec_size: Scalefactor A/B vector size.
- mma_tiler_mn: The (M, N) shape of the MMA instruction tiler.
2. Cluster Shape:
- cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster.
3. TMA Prefetch:
- prefetch_dist: Prefetch distance for TMA operations.
None = use num_ab_stage (default), 0 = disable prefetch, >0 = explicit distance.
:param sf_vec_size: Scalefactor vector size.
:type sf_vec_size: int
:param mma_tiler_mn: Tuple (M, N) shape of the MMA instruction.
:type mma_tiler_mn: Tuple[int, int]
:param cluster_shape_mn: Tuple (ClusterM, ClusterN) shape of the cluster.
:type cluster_shape_mn: Tuple[int, int]
:param prefetch_dist: Prefetch distance for TMA operations (None=auto, 0=disable, >0=explicit).
:type prefetch_dist: Union[int, None]
"""
mma_tiler_mn = config_map[problem_size][0]
cluster_shape_mn = config_map[problem_size][1]
self.appro = config_map[problem_size][3]
self.acc_dtype = cutlass.Float32
self.sf_vec_size = sf_vec_size
self.use_2cta_instrs = mma_tiler_mn[0] == 256
self.cluster_shape_mn = cluster_shape_mn
# K dimension is deferred in _setup_attributes
self.mma_tiler = (*mma_tiler_mn, 1)
# Prefetch configuration: None=auto (num_ab_stage), 0=disable, >0=explicit distance
self.prefetch_dist_param = config_map[problem_size][2]
self.cta_group = (
tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
)
self.occupancy = 1
# Set specialized warp ids
self.epilog_warp_id = (
0,
1,
2,
3,
)
self.mma_warp_id = 4
self.tma_warp_id = 5
self.threads_per_cta = 32 * len(
(self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)
)
# Set barrier id for epilogue sync and tmem ptr sync
self.epilog_sync_barrier = pipeline.NamedBarrier(
barrier_id=1,
num_threads=32 * len(self.epilog_warp_id),
)
self.tmem_alloc_barrier = pipeline.NamedBarrier(
barrier_id=2,
num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)),
)
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
SM100_TMEM_CAPACITY_COLUMNS = 512
self.num_tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS
def _setup_attributes(self):
"""Set up configurations that are dependent on GEMM inputs
This method configures various attributes based on the input tensor properties
(data types, leading dimensions) and kernel settings:
- Configuring tiled MMA
- Computing MMA/cluster/tile shapes
- Computing cluster layout
- Computing multicast CTAs for A/B/SFA/SFB
- Computing epilogue subtile
- Setting up A/B/SFA/SFB/C stage counts in shared memory
- Computing A/B/SFA/SFB/C shared memory layout
"""
# Compute mma instruction shapes
# (MMA_Tile_Shape_M, MMA_Tile_Shape_N, MMA_Inst_Shape_K)
self.mma_inst_shape_mn = (
self.mma_tiler[0],
self.mma_tiler[1],
)
# (CTA_Tile_Shape_M, Round_Up(MMA_Tile_Shape_N, 128), MMA_Inst_Shape_K)
self.mma_inst_shape_mn_sfb = (
self.mma_inst_shape_mn[0] // (2 if self.use_2cta_instrs else 1),
cute.round_up(self.mma_inst_shape_mn[1], 128),
)
tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.a_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
self.cta_group,
self.mma_inst_shape_mn,
)
tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.a_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
cute.nvgpu.tcgen05.CtaGroup.ONE,
self.mma_inst_shape_mn_sfb,
)
# Compute mma/cluster/tile shapes
mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
mma_inst_tile_k = 4
self.mma_tiler = (
self.mma_inst_shape_mn[0],
self.mma_inst_shape_mn[1],
mma_inst_shape_k * mma_inst_tile_k,
)
self.mma_tiler_sfb = (
self.mma_inst_shape_mn_sfb[0],
self.mma_inst_shape_mn_sfb[1],
mma_inst_shape_k * mma_inst_tile_k,
)
self.cta_tile_shape_mnk = (
self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape),
self.mma_tiler[1],
self.mma_tiler[2],
)
self.cta_tile_shape_mnk_sfb = (
self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape),
self.mma_tiler_sfb[1],
self.mma_tiler_sfb[2],
)
# Compute cluster layout
self.cluster_layout_vmnk = cute.tiled_divide(
cute.make_layout((*self.cluster_shape_mn, 1)),
(tiled_mma.thr_id.shape,),
)
self.cluster_layout_sfb_vmnk = cute.tiled_divide(
cute.make_layout((*self.cluster_shape_mn, 1)),
(tiled_mma_sfb.thr_id.shape,),
)
# Compute number of multicast CTAs for A/B
self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2])
self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1])
self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1])
self.is_a_mcast = self.num_mcast_ctas_a > 1
self.is_b_mcast = self.num_mcast_ctas_b > 1
self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1
# Compute epilogue subtile
self.epi_tile = sm100_utils.compute_epilogue_tile_shape(
self.cta_tile_shape_mnk,
self.use_2cta_instrs,
self.c_layout,
self.c_dtype,
)
self.epi_tile_n = cute.size(self.epi_tile[1])
# Setup A/B/C stage count in shared memory and ACC stage count in tensor memory
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,
)
# Compute A/B/SFA/SFB/C shared memory layout
self.a_smem_layout_staged = sm100_utils.make_smem_layout_a(
tiled_mma,
self.mma_tiler,
self.a_dtype,
self.num_ab_stage,
)
self.b_smem_layout_staged = sm100_utils.make_smem_layout_b(
tiled_mma,
self.mma_tiler,
self.b_dtype,
self.num_ab_stage,
)
self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa(
tiled_mma,
self.mma_tiler,
self.sf_vec_size,
self.num_ab_stage,
)
self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb(
tiled_mma,
self.mma_tiler,
self.sf_vec_size,
self.num_ab_stage,
)
self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi(
self.c_dtype,
self.c_layout,
self.epi_tile,
self.num_c_stage,
)
# Overlap and double buffer accumulator when num_acc_stage == 1 for cta_tile_n = 256 case
# self.overlapping_accum = self.num_acc_stage == 1
self.overlapping_accum = False
# Compute number of TMEM columns for SFA/SFB/Accumulator
sf_atom_mn = 32
self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k
self.num_sfb_tmem_cols = (self.cta_tile_shape_mnk_sfb[1] // sf_atom_mn) * mma_inst_tile_k
self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols
self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] * self.num_acc_stage if not self.overlapping_accum else self.cta_tile_shape_mnk[1] * 2 - self.num_sf_tmem_cols
# Only when overlapping_accum is enabled, we need to release accumulator buffer early in epilogue
self.iter_acc_early_release_in_epilogue = self.num_sf_tmem_cols // self.epi_tile_n
# Set prefetch distance for both initial and rolling prefetch (unified control)
# None = use num_ab_stage (default), 0 = disable prefetch, >0 = explicit distance
if self.prefetch_dist_param is None:
self.prefetch_dist = self.num_ab_stage
else:
self.prefetch_dist = self.prefetch_dist_param
# Check if prefetch is enabled (prefetch_dist > 0)
self.prefetch_enabled = self.prefetch_dist > 0
@cute.jit
def __call__(
self,
a_ptr: cute.Pointer,
b1_ptr: cute.Pointer,
b2_ptr: cute.Pointer,
sfa_ptr: cute.Pointer,
sfb1_ptr: cute.Pointer,
sfb2_ptr: cute.Pointer,
c_ptr: cute.Pointer,
problem_size: cutlass.Constexpr,
# epilogue_op: cutlass.Constexpr = lambda x: x,
epilogue_op: cutlass.Constexpr = lambda x: 0.5*x * (1.0 + cute.math.tanh(0.5*x, fastmath=True)),
):
m, n, k, l = problem_size
# Setup attributes that depend on gemm inputs
a_tensor = cute.make_tensor(
a_ptr,
cute.make_layout(
(m, k, l),
stride=(k, 1, m * k),
),
)
b1_tensor = cute.make_tensor(
b1_ptr,
cute.make_layout(
(n, k, l),
stride=(k, 1, n * k),
),
)
b2_tensor = cute.make_tensor(
b2_ptr,
cute.make_layout(
(n, k, l),
stride=(k, 1, n * k),
),
)
c_tensor = cute.make_tensor(
c_ptr, cute.make_layout((m, n, l), stride=(n, 1, m * n))
)
"""Execute the GEMM operation in steps:
- Setup static attributes before smem/grid/tma computation
- Setup TMA load/store atoms and tensors
- Compute grid size with regard to hardware constraints
- Define shared storage for kernel
- Launch the kernel synchronously
:param a_tensor: Input tensor A
:type a_tensor: cute.Tensor
:param b_tensor: Input tensor B
:type b_tensor: cute.Tensor
:param sfa_tensor: Scale factor tensor A
:type sfa_tensor: cute.Tensor
:param sfb_tensor: Scale factor tensor B
:type sfb_tensor: cute.Tensor
:param c_tensor: Output tensor C
:type c_tensor: cute.Tensor
:param max_active_clusters: Maximum number of active clusters
:type max_active_clusters: cutlass.Constexpr
:param epilogue_op: Optional elementwise lambda function to apply to the output tensor
:type epilogue_op: cutlass.Constexpr
:raises TypeError: If input data types are incompatible with the MMA instruction.
"""
# Setup static attributes before smem/grid/tma computation
self.a_dtype: Type[cutlass.Numeric] = a_tensor.element_type
self.b_dtype: Type[cutlass.Numeric] = b1_tensor.element_type
self.sf_dtype: Type[cutlass.Numeric] = sf_dtype
self.c_dtype: Type[cutlass.Numeric] = c_tensor.element_type
self.a_major_mode = utils.LayoutEnum.from_tensor(a_tensor).mma_major_mode()
self.b_major_mode = utils.LayoutEnum.from_tensor(b1_tensor).mma_major_mode()
self.c_layout = utils.LayoutEnum.from_tensor(c_tensor)
# Check if input data types are compatible with MMA instruction
if cutlass.const_expr(self.a_dtype != self.b_dtype):
raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}")
# Setup attributes that dependent on gemm inputs
self._setup_attributes()
# Setup sfa/sfb tensor by filling A/B tensor to scale factor atom layout
# ((Atom_M, Rest_M),(Atom_K, Rest_K),RestL)
sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(
a_tensor.shape, self.sf_vec_size
)
sfa_tensor = cute.make_tensor(sfa_ptr, sfa_layout)
# ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL)
sfb1_layout = blockscaled_utils.tile_atom_to_shape_SF(
b1_tensor.shape, self.sf_vec_size
)
sfb1_tensor = cute.make_tensor(sfb1_ptr, sfb1_layout)
sfb2_layout = blockscaled_utils.tile_atom_to_shape_SF(
b2_tensor.shape, self.sf_vec_size
)
sfb2_tensor = cute.make_tensor(sfb2_ptr, sfb2_layout)
tiled_mma = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.a_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
self.cta_group,
self.mma_inst_shape_mn,
)
tiled_mma_sfb = sm100_utils.make_blockscaled_trivial_tiled_mma(
self.a_dtype,
self.a_major_mode,
self.b_major_mode,
self.sf_dtype,
self.sf_vec_size,
cute.nvgpu.tcgen05.CtaGroup.ONE,
self.mma_inst_shape_mn_sfb,
)
atom_thr_size = cute.size(tiled_mma.thr_id.shape)
# Setup TMA load for A
a_op = sm100_utils.cluster_shape_to_tma_atom_A(
self.cluster_shape_mn, tiled_mma.thr_id
)
a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0))
tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A(
a_op,
a_tensor,
a_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
)
# Setup TMA load for B
b_op = sm100_utils.cluster_shape_to_tma_atom_B(
self.cluster_shape_mn, tiled_mma.thr_id
)
b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0))
tma_atom_b1, tma_tensor_b1 = cute.nvgpu.make_tiled_tma_atom_B(
b_op,
b1_tensor,
b_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
)
tma_atom_b2, tma_tensor_b2 = cute.nvgpu.make_tiled_tma_atom_B(
b_op,
b2_tensor,
b_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
)
# Setup TMA load for SFA
sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(
self.cluster_shape_mn, tiled_mma.thr_id
)
sfa_smem_layout = cute.slice_(
self.sfa_smem_layout_staged, (None, None, None, 0)
)
tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A(
sfa_op,
sfa_tensor,
sfa_smem_layout,
self.mma_tiler,
tiled_mma,
self.cluster_layout_vmnk.shape,
internal_type=cutlass.Int16,
)
# Setup TMA load for SFB
sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(
self.cluster_shape_mn, tiled_mma.thr_id
)
sfb_smem_layout = cute.slice_(
self.sfb_smem_layout_staged, (None, None, None, 0)
)
tma_atom_sfb1, tma_tensor_sfb1 = cute.nvgpu.make_tiled_tma_atom_B(
sfb_op,
sfb1_tensor,
sfb_smem_layout,
self.mma_tiler_sfb,
tiled_mma_sfb,
self.cluster_layout_sfb_vmnk.shape,
internal_type=cutlass.Int16,
)
tma_atom_sfb2, tma_tensor_sfb2 = cute.nvgpu.make_tiled_tma_atom_B(
sfb_op,
sfb2_tensor,
sfb_smem_layout,
self.mma_tiler_sfb,
tiled_mma_sfb,
self.cluster_layout_sfb_vmnk.shape,
internal_type=cutlass.Int16,
)
if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
x = tma_tensor_sfb1.stride[0][1]
y = cute.ceil_div(tma_tensor_sfb1.shape[0][1], 4)
new_shape = (
(
tma_tensor_sfb1.shape[0][0],
((2, 2), y)
),
tma_tensor_sfb1.shape[1],
tma_tensor_sfb1.shape[2]
)
# Use right multiplication for ScaledBasis (3 * x instead of x * 3)
x_times_3 = 3 * x
new_stride = (
(
tma_tensor_sfb1.stride[0][0],
((x, x), x_times_3)
),
tma_tensor_sfb1.stride[1],
tma_tensor_sfb1.stride[2]
)
tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride)
tma_tensor_sfb1 = cute.make_tensor(tma_tensor_sfb1.iterator, tma_tensor_sfb_new_layout)
tma_tensor_sfb2 = cute.make_tensor(tma_tensor_sfb2.iterator, tma_tensor_sfb_new_layout)
a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout)
b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout)
sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout)
sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout)
self.num_tma_load_bytes = (
a_copy_size + b_copy_size * 2 + sfa_copy_size + sfb_copy_size * 2
) * atom_thr_size
# Setup TMA store for C
epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0))
tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom(
cpasync.CopyBulkTensorTileS2GOp(),
c_tensor,
epi_smem_layout,
self.epi_tile,
)
# Compute grid size
self.tile_sched_params, grid = self._compute_grid(
c_tensor,
self.cta_tile_shape_mnk,
self.cluster_shape_mn,
max_active_clusters,
)
self.buffer_align_bytes = 1024
# Define shared storage for kernel
@cute.struct
class SharedStorage:
ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
ab_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage]
acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
acc_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage]
tmem_dealloc_mbar_ptr: cutlass.Int64
tmem_holding_buf: cutlass.Int32
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sC: cute.struct.Align[
cute.struct.MemRange[
self.c_dtype,
cute.cosize(self.c_smem_layout_staged.outer),
],
self.buffer_align_bytes,
]
# (MMA, MMA_M, MMA_K, STAGE)
sA: cute.struct.Align[
cute.struct.MemRange[
self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)
],
self.buffer_align_bytes,
]
# (MMA, MMA_N, MMA_K, STAGE)
sB1: cute.struct.Align[
cute.struct.MemRange[
self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)
],
self.buffer_align_bytes,
]
sB2: cute.struct.Align[
cute.struct.MemRange[
self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)
],
self.buffer_align_bytes,
]
# (MMA, MMA_M, MMA_K, STAGE)
sSFA: cute.struct.Align[
cute.struct.MemRange[
self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)
],
self.buffer_align_bytes,
]
# (MMA, MMA_N, MMA_K, STAGE)
sSFB1: cute.struct.Align[
cute.struct.MemRange[
self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)
],
self.buffer_align_bytes,
]
sSFB2: cute.struct.Align[
cute.struct.MemRange[
self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)
],
self.buffer_align_bytes,
]
self.shared_storage = SharedStorage
# Launch the kernel synchronously
self.kernel(
tiled_mma,
tiled_mma_sfb,
tma_atom_a,
tma_tensor_a,
tma_atom_b1,
tma_tensor_b1,
tma_atom_b2,
tma_tensor_b2,
tma_atom_sfa,
tma_tensor_sfa,
tma_atom_sfb1,
tma_tensor_sfb1,
tma_atom_sfb2,
tma_tensor_sfb2,
tma_atom_c,
tma_tensor_c,
self.cluster_layout_vmnk,
self.cluster_layout_sfb_vmnk,
self.a_smem_layout_staged,
self.b_smem_layout_staged,
self.sfa_smem_layout_staged,
self.sfb_smem_layout_staged,
self.c_smem_layout_staged,
self.epi_tile,
self.tile_sched_params,
epilogue_op,
).launch(
grid=grid,
block=[self.threads_per_cta, 1, 1],
cluster=(*self.cluster_shape_mn, 1),
min_blocks_per_mp=1,
)
return
# GPU device kernel
@cute.kernel
def kernel(
self,
tiled_mma: cute.TiledMma,
tiled_mma_sfb: cute.TiledMma,
tma_atom_a: cute.CopyAtom,
mA_mkl: cute.Tensor,
tma_atom_b1: cute.CopyAtom,
mB1_nkl: cute.Tensor,
tma_atom_b2: cute.CopyAtom,
mB2_nkl: cute.Tensor,
tma_atom_sfa: cute.CopyAtom,
mSFA_mkl: cute.Tensor,
tma_atom_sfb1: cute.CopyAtom,
mSFB1_nkl: cute.Tensor,
tma_atom_sfb2: cute.CopyAtom,
mSFB2_nkl: cute.Tensor,
tma_atom_c: cute.CopyAtom,
mC_mnl: cute.Tensor,
cluster_layout_vmnk: cute.Layout,
cluster_layout_sfb_vmnk: cute.Layout,
a_smem_layout_staged: cute.ComposedLayout,
b_smem_layout_staged: cute.ComposedLayout,
sfa_smem_layout_staged: cute.Layout,
sfb_smem_layout_staged: cute.Layout,
c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout],
epi_tile: cute.Tile,
tile_sched_params: utils.PersistentTileSchedulerParams,
epilogue_op: cutlass.Constexpr,
):
"""
GPU device kernel performing the Persistent batched GEMM computation.
"""
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
#
# Prefetch tma desc
#
if warp_idx == self.tma_warp_id:
cpasync.prefetch_descriptor(tma_atom_a)
cpasync.prefetch_descriptor(tma_atom_b1)
cpasync.prefetch_descriptor(tma_atom_b2)
cpasync.prefetch_descriptor(tma_atom_sfa)
cpasync.prefetch_descriptor(tma_atom_sfb1)
cpasync.prefetch_descriptor(tma_atom_sfb2)
cpasync.prefetch_descriptor(tma_atom_c)
use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2
#
# Setup cta/thread coordinates
#
# Coords inside cluster
bidx, bidy, bidz = cute.arch.block_idx()
mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape)
is_leader_cta = mma_tile_coord_v == 0
cta_rank_in_cluster = cute.arch.make_warp_uniform(
cute.arch.block_idx_in_cluster()
)
block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(
cta_rank_in_cluster
)
block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord(
cta_rank_in_cluster
)
# Coord inside cta
tidx, _, _ = cute.arch.thread_idx()
#
# Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier
#
smem = utils.SmemAllocator()
storage = smem.allocate(self.shared_storage)
# Initialize mainloop ab_pipeline (barrier) and states
ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1
ab_pipeline_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, num_tma_producer
)
ab_pipeline = pipeline.PipelineTmaUmma.create(
barrier_storage=storage.ab_full_mbar_ptr.data_ptr(),
num_stages=self.num_ab_stage,
producer_group=ab_pipeline_producer_group,
consumer_group=ab_pipeline_consumer_group,
tx_count=self.num_tma_load_bytes,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Initialize acc_pipeline (barrier) and states
acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread)
num_acc_consumer_threads = len(self.epilog_warp_id) * (
2 if use_2cta_instrs else 1
)
acc_pipeline_consumer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread, num_acc_consumer_threads
)
acc_pipeline = pipeline.PipelineUmmaAsync.create(
barrier_storage=storage.acc_full_mbar_ptr.data_ptr(),
num_stages=self.num_acc_stage,
producer_group=acc_pipeline_producer_group,
consumer_group=acc_pipeline_consumer_group,
cta_layout_vmnk=cluster_layout_vmnk,
defer_sync=True,
)
# Tensor memory dealloc barrier init
tmem = utils.TmemAllocator(
storage.tmem_holding_buf,
barrier_for_retrieve=self.tmem_alloc_barrier,
allocator_warp_id=self.epilog_warp_id[0],
is_two_cta=use_2cta_instrs,
two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr,
)
# Cluster arrive after barrier init
# pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True)
#
# Setup smem tensor A/B/SFA/SFB/C
#
# (EPI_TILE_M, EPI_TILE_N, STAGE)
sC = storage.sC.get_tensor(
c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner
)
# (MMA, MMA_M, MMA_K, STAGE)
sA = storage.sA.get_tensor(
a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner
)
# (MMA, MMA_N, MMA_K, STAGE)
sB1 = storage.sB1.get_tensor(
b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner
)
sB2 = storage.sB2.get_tensor(
b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner
)
# (MMA, MMA_M, MMA_K, STAGE)
sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged)
# (MMA, MMA_N, MMA_K, STAGE)
sSFB1 = storage.sSFB1.get_tensor(sfb_smem_layout_staged)
sSFB2 = storage.sSFB2.get_tensor(sfb_smem_layout_staged)
#
# Compute multicast mask for A/B/SFA/SFB buffer full
#
a_full_mcast_mask = None
b_full_mcast_mask = None
sfa_full_mcast_mask = None
sfb_full_mcast_mask = None
if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs):
a_full_mcast_mask = cpasync.create_tma_multicast_mask(
cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
)
b_full_mcast_mask = cpasync.create_tma_multicast_mask(
cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1
)
sfa_full_mcast_mask = cpasync.create_tma_multicast_mask(
cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2
)
sfb_full_mcast_mask = cpasync.create_tma_multicast_mask(
cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1
)
#
# Local_tile partition global tensors
#
# (bM, bK, RestM, RestK, RestL)
gA_mkl = cute.local_tile(
mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
# (bN, bK, RestN, RestK, RestL)
gB1_nkl = cute.local_tile(
mB1_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
gB2_nkl = cute.local_tile(
mB2_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)
)
# (bM, bK, RestM, RestK, RestL)
gSFA_mkl = cute.local_tile(
mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)
)
# (bN, bK, RestN, RestK, RestL)
gSFB1_nkl = cute.local_tile(
mSFB1_nkl,
cute.slice_(self.mma_tiler_sfb, (0, None, None)),
(None, None, None),
)
gSFB2_nkl = cute.local_tile(
mSFB2_nkl,
cute.slice_(self.mma_tiler_sfb, (0, None, None)),
(None, None, None),
)
# (bM, bN, RestM, RestN, RestL)
gC_mnl = cute.local_tile(
mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
)
k_tile_cnt = cute.size(gA_mkl, mode=[3])
#
# Partition global tensor for TiledMMA_A/B/C
#
thr_mma = tiled_mma.get_slice(mma_tile_coord_v)
thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v)
# (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
tCgA = thr_mma.partition_A(gA_mkl)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgB1 = thr_mma.partition_B(gB1_nkl)
tCgB2 = thr_mma.partition_B(gB2_nkl)
# (MMA, MMA_M, MMA_K, RestM, RestK, RestL)
tCgSFA = thr_mma.partition_A(gSFA_mkl)
# (MMA, MMA_N, MMA_K, RestN, RestK, RestL)
tCgSFB1 = thr_mma_sfb.partition_B(gSFB1_nkl)
tCgSFB2 = thr_mma_sfb.partition_B(gSFB2_nkl)
# (MMA, MMA_M, MMA_N, RestM, RestN, RestL)
tCgC = thr_mma.partition_C(gC_mnl)
#
# Partition global/shared tensor for TMA load A/B
#
# TMA load A partition_S/D
a_cta_layout = cute.make_layout(
cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape
)
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestM, RestK, RestL)
tAsA, tAgA = cpasync.tma_partition(
tma_atom_a,
block_in_cluster_coord_vmnk[2],
a_cta_layout,
cute.group_modes(sA, 0, 3),
cute.group_modes(tCgA, 0, 3),
)
# TMA load B partition_S/D
b_cta_layout = cute.make_layout(
cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
)
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestN, RestK, RestL)
tBsB1, tBgB1 = cpasync.tma_partition(
tma_atom_b1,
block_in_cluster_coord_vmnk[1],
b_cta_layout,
cute.group_modes(sB1, 0, 3),
cute.group_modes(tCgB1, 0, 3),
)
tBsB2, tBgB2 = cpasync.tma_partition(
tma_atom_b2,
block_in_cluster_coord_vmnk[1],
b_cta_layout,
cute.group_modes(sB2, 0, 3),
cute.group_modes(tCgB2, 0, 3),
)
# TMA load SFA partition_S/D
sfa_cta_layout = a_cta_layout
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestM, RestK, RestL)
tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition(
tma_atom_sfa,
block_in_cluster_coord_vmnk[2],
sfa_cta_layout,
cute.group_modes(sSFA, 0, 3),
cute.group_modes(tCgSFA, 0, 3),
)
tAsSFA = cute.filter_zeros(tAsSFA)
tAgSFA = cute.filter_zeros(tAgSFA)
# TMA load SFB partition_S/D
sfb_cta_layout = cute.make_layout(
cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape
)
# ((atom_v, rest_v), STAGE)
# ((atom_v, rest_v), RestN, RestK, RestL)
tBsSFB1, tBgSFB1 = cute.nvgpu.cpasync.tma_partition(
tma_atom_sfb1,
block_in_cluster_coord_sfb_vmnk[1],
sfb_cta_layout,
cute.group_modes(sSFB1, 0, 3),
cute.group_modes(tCgSFB1, 0, 3),
)
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)
#
# Partition shared/tensor memory tensor for TiledMMA_A/B/C
#
# (MMA, MMA_M, MMA_K, STAGE)
tCrA = tiled_mma.make_fragment_A(sA)
# (MMA, MMA_N, MMA_K, STAGE)
tCrB1 = tiled_mma.make_fragment_B(sB1)
tCrB2 = tiled_mma.make_fragment_B(sB2)
# (MMA, MMA_M, MMA_N)
acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2])
if cutlass.const_expr(self.overlapping_accum):
num_acc_stage_overlapped = 2
tCtAcc_fake = tiled_mma.make_fragment_C(
cute.append(acc_shape, num_acc_stage_overlapped)
)
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_fake = cute.make_tensor(
tCtAcc_fake.iterator,
cute.make_layout(
tCtAcc_fake.shape,
stride = (
tCtAcc_fake.stride[0],
tCtAcc_fake.stride[1],
tCtAcc_fake.stride[2],
(256 - self.num_sf_tmem_cols) * tCtAcc_fake.stride[0][1]
)
)
)
else:
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_fake = tiled_mma.make_fragment_C(
cute.append(acc_shape, self.num_acc_stage)
)
#
# Cluster wait before tensor memory alloc
#
# pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn)
#
# Specialized TMA load warp
#
if warp_idx == self.tma_warp_id:
#
# Persistent tile scheduling loop
#
tile_sched = utils.StaticPersistentTileScheduler.create(
tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
)
work_tile = tile_sched.initial_work_tile_info()
ab_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_ab_stage
)
while work_tile.is_valid_tile:
# Get tile coord from tile scheduler
cur_tile_coord = work_tile.tile_idx
mma_tile_coord_mnl = (
cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
cur_tile_coord[1],
cur_tile_coord[2],
)
#
# Slice to per mma tile index
#
# ((atom_v, rest_v), RestK)
tAgA_slice = tAgA[
(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])
]
# ((atom_v, rest_v), RestK)
tBgB1_slice = tBgB1[
(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])
]
tBgB2_slice = tBgB2[
(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])
]
# ((atom_v, rest_v), RestK)
tAgSFA_slice = tAgSFA[
(None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])
]
slice_n = mma_tile_coord_mnl[1]
if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
slice_n = mma_tile_coord_mnl[1] // 2
# ((atom_v, rest_v), RestK)
tBgSFB1_slice = tBgSFB1[
(None, slice_n, None, mma_tile_coord_mnl[2])
]
tBgSFB2_slice = tBgSFB2[
(None, slice_n, None, mma_tile_coord_mnl[2])
]
#
# Prefetch: Initial batch of prefetches to prime the pipeline
#
if self.prefetch_enabled:
for pf_k_tile in cutlass.range(
0, min(self.prefetch_dist, k_tile_cnt), unroll=1
):
cute.prefetch(
tma_atom_a,
tAgA_slice[(None, pf_k_tile)],
)
cute.prefetch(
tma_atom_b1,
tBgB1_slice[(None, pf_k_tile)],
)
cute.prefetch(
tma_atom_b2,
tBgB2_slice[(None, pf_k_tile)],
)
cute.prefetch(
tma_atom_sfa,
tAgSFA_slice[(None, pf_k_tile)],
)
cute.prefetch(
tma_atom_sfb1,
tBgSFB1_slice[(None, pf_k_tile)],
)
cute.prefetch(
tma_atom_sfb2,
tBgSFB2_slice[(None, pf_k_tile)],
)
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt
ab_producer_state.reset_count()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < k_tile_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(
ab_producer_state
)
#
# Tma load loop
#
for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):
# Conditionally wait for AB buffer empty
ab_pipeline.producer_acquire(
ab_producer_state, peek_ab_empty_status
)
# TMA load A/B/SFA/SFB
cute.copy(
tma_atom_a,
tAgA_slice[(None, ab_producer_state.count)],
tAsA[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=a_full_mcast_mask,
)
cute.copy(
tma_atom_b1,
tBgB1_slice[(None, ab_producer_state.count)],
tBsB1[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=b_full_mcast_mask,
)
cute.copy(
tma_atom_b2,
tBgB2_slice[(None, ab_producer_state.count)],
tBsB2[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=b_full_mcast_mask,
)
cute.copy(
tma_atom_sfa,
tAgSFA_slice[(None, ab_producer_state.count)],
tAsSFA[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=sfa_full_mcast_mask,
)
cute.copy(
tma_atom_sfb1,
tBgSFB1_slice[(None, ab_producer_state.count)],
tBsSFB1[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=sfb_full_mcast_mask,
)
cute.copy(
tma_atom_sfb2,
tBgSFB2_slice[(None, ab_producer_state.count)],
tBsSFB2[(None, ab_producer_state.index)],
tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),
mcast_mask=sfb_full_mcast_mask,
)
# Prefetch: Rolling prefetch for next tiles
if self.prefetch_enabled:
if k_tile < k_tile_cnt - self.prefetch_dist:
future_k_tile = ab_producer_state.count + self.prefetch_dist
cute.prefetch(
tma_atom_a,
tAgA_slice[(None, future_k_tile)],
)
cute.prefetch(
tma_atom_b1,
tBgB1_slice[(None, future_k_tile)],
)
cute.prefetch(
tma_atom_b2,
tBgB2_slice[(None, future_k_tile)],
)
cute.prefetch(
tma_atom_sfa,
tAgSFA_slice[(None, future_k_tile)],
)
cute.prefetch(
tma_atom_sfb1,
tBgSFB1_slice[(None, future_k_tile)],
)
cute.prefetch(
tma_atom_sfb2,
tBgSFB2_slice[(None, future_k_tile)],
)
# Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1
ab_producer_state.advance()
peek_ab_empty_status = cutlass.Boolean(1)
if ab_producer_state.count < k_tile_cnt:
peek_ab_empty_status = ab_pipeline.producer_try_acquire(
ab_producer_state
)
#
# Advance to next tile
#
tile_sched.advance_to_next_work()
work_tile = tile_sched.get_current_work()
#
# Wait A/B buffer empty
#
ab_pipeline.producer_tail(ab_producer_state)
#
# Specialized MMA warp
#
if warp_idx == self.mma_warp_id:
#
# Bar sync for retrieve tensor memory ptr from shared mem
#
tmem.wait_for_alloc()
#
# Retrieving tensor memory ptr and make accumulator/SFA/SFB tensor
#
acc1_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
tCtAcc1_base = cute.make_tensor(acc1_tmem_ptr, tCtAcc_fake.layout)
acc2_tmem_ptr = cute.recast_ptr(
acc1_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base),
dtype=self.acc_dtype,
)
tCtAcc2_base = cute.make_tensor(acc2_tmem_ptr, tCtAcc_fake.layout)
# Make SFA tmem tensor
sfa_tmem_ptr = cute.recast_ptr(
acc1_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2,
dtype=self.sf_dtype,
)
# (MMA, MMA_M, MMA_K)
tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa(
tiled_mma,
self.mma_tiler,
self.sf_vec_size,
cute.slice_(sfa_smem_layout_staged, (None, None, None, 0)),
)
tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout)
# Make SFB tmem tensor
# (MMA, MMA_N, MMA_K)
tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb(
tiled_mma,
self.mma_tiler,
self.sf_vec_size,
cute.slice_(sfb_smem_layout_staged, (None, None, None, 0)),
)
# Make SFB tmem tensor
sfb1_tmem_ptr = cute.recast_ptr(
acc1_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA),
dtype=self.sf_dtype,
)
tCtSFB1 = cute.make_tensor(sfb1_tmem_ptr, tCtSFB_layout)
sfb2_tmem_ptr = cute.recast_ptr(
acc1_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2
+ 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)
#
# Partition for S2T copy of SFA/SFB
#
(
tiled_copy_s2t_sfa,
tCsSFA_compact_s2t,
tCtSFA_compact_s2t,
) = self.mainloop_s2t_copy_and_partition(sSFA, tCtSFA)
(
tiled_copy_s2t_sfb1,
tCsSFB1_compact_s2t,
tCtSFB1_compact_s2t,
) = self.mainloop_s2t_copy_and_partition(sSFB1, tCtSFB1)
(
tiled_copy_s2t_sfb2,
tCsSFB2_compact_s2t,
tCtSFB2_compact_s2t,
) = self.mainloop_s2t_copy_and_partition(sSFB2, tCtSFB2)
#
# Persistent tile scheduling loop
#
tile_sched = utils.StaticPersistentTileScheduler.create(
tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
)
work_tile = tile_sched.initial_work_tile_info()
ab_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_ab_stage
)
acc_producer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Producer, self.num_acc_stage
)
while work_tile.is_valid_tile:
# Get tile coord from tile scheduler
cur_tile_coord = work_tile.tile_idx
mma_tile_coord_mnl = (
cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
cur_tile_coord[1],
cur_tile_coord[2],
)
# Get accumulator stage index
if cutlass.const_expr(self.overlapping_accum):
acc_stage_index = acc_producer_state.phase ^ 1
else:
acc_stage_index = acc_producer_state.index
# Set tensor memory buffer for current tile
# (MMA, MMA_M, MMA_N)
tCtAcc1 = tCtAcc1_base[(None, None, None, acc_producer_state.index)]
tCtAcc2 = tCtAcc2_base[(None, None, None, acc_producer_state.index)]
# Peek (try_wait) AB buffer full for k_tile = 0
ab_consumer_state.reset_count()
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < k_tile_cnt and is_leader_cta:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
)
#
# Wait for accumulator buffer empty
#
if is_leader_cta:
acc_pipeline.producer_acquire(acc_producer_state)
tCtSFB1_mma = tCtSFB1
tCtSFB2_mma = tCtSFB2
if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):
# If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words (ignores first 64 columns of SFB)
offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0)
shifted_ptr1 = cute.recast_ptr(
acc1_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+ offset,
dtype=self.sf_dtype,
)
shifted_ptr2 = cute.recast_ptr(
acc1_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFB1)
+ offset,
dtype=self.sf_dtype,
)
tCtSFB1_mma = cute.make_tensor(shifted_ptr1, tCtSFB_layout)
tCtSFB2_mma = cute.make_tensor(shifted_ptr2, tCtSFB_layout)
elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):
# Move in increments of 64 columns of SFB
offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)
shifted_ptr1 = cute.recast_ptr(
acc1_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+ offset,
dtype=self.sf_dtype,
)
shifted_ptr2 = cute.recast_ptr(
acc1_tmem_ptr
+ tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2
+ tcgen05.find_tmem_tensor_col_offset(tCtSFA)
+ tcgen05.find_tmem_tensor_col_offset(tCtSFB1)
+ offset,
dtype=self.sf_dtype,
)
tCtSFB1_mma = cute.make_tensor(shifted_ptr1, tCtSFB_layout)
tCtSFB2_mma = cute.make_tensor(shifted_ptr2, tCtSFB_layout)
#
# Reset the ACCUMULATE field for each tile
#
tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
#
# Mma mainloop
#
for k_tile in range(k_tile_cnt):
if is_leader_cta:
# Conditionally wait for AB buffer full
ab_pipeline.consumer_wait(
ab_consumer_state, peek_ab_full_status
)
# Copy SFA/SFB from smem to tmem
s2t_stage_coord = (
None,
None,
None,
None,
ab_consumer_state.index,
)
tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]
tCsSFB1_compact_s2t_staged = tCsSFB1_compact_s2t[s2t_stage_coord]
tCsSFB2_compact_s2t_staged = tCsSFB2_compact_s2t[s2t_stage_coord]
cute.copy(
tiled_copy_s2t_sfa,
tCsSFA_compact_s2t_staged,
tCtSFA_compact_s2t,
)
cute.copy(
tiled_copy_s2t_sfb1,
tCsSFB1_compact_s2t_staged,
tCtSFB1_compact_s2t,
)
cute.copy(
tiled_copy_s2t_sfb2,
tCsSFB2_compact_s2t_staged,
tCtSFB2_compact_s2t,
)
# tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB
num_kblocks = cute.size(tCrA, mode=[2])
for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):
kblock_coord = (
None,
None,
kblock_idx,
ab_consumer_state.index,
)
# Set SFA/SFB tensor to tiled_mma
sf_kblock_coord = (None, None, kblock_idx)
tiled_mma.set(
tcgen05.Field.SFA,
tCtSFA[sf_kblock_coord].iterator,
)
tiled_mma.set(
tcgen05.Field.SFB,
tCtSFB1_mma[sf_kblock_coord].iterator,
)
cute.gemm(
tiled_mma,
tCtAcc1,
tCrA[kblock_coord],
tCrB1[kblock_coord],
tCtAcc1,
)
tiled_mma.set(
tcgen05.Field.SFB,
tCtSFB2_mma[sf_kblock_coord].iterator,
)
cute.gemm(
tiled_mma,
tCtAcc2,
tCrA[kblock_coord],
tCrB2[kblock_coord],
tCtAcc2,
)
# Enable accumulate on tCtAcc after first kblock
tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
# Async arrive AB buffer empty
ab_pipeline.consumer_release(ab_consumer_state)
# Peek (try_wait) AB buffer full for k_tile = k_tile + 1
ab_consumer_state.advance()
peek_ab_full_status = cutlass.Boolean(1)
if ab_consumer_state.count < k_tile_cnt:
if is_leader_cta:
peek_ab_full_status = ab_pipeline.consumer_try_wait(
ab_consumer_state
)
#
# Async arrive accumulator buffer full
#
if is_leader_cta:
acc_pipeline.producer_commit(acc_producer_state)
acc_producer_state.advance()
#
# Advance to next tile
#
tile_sched.advance_to_next_work()
work_tile = tile_sched.get_current_work()
#
# Wait for accumulator buffer empty
#
acc_pipeline.producer_tail(acc_producer_state)
#
# Specialized epilogue warps
#
if warp_idx < self.mma_warp_id:
#
# Alloc tensor memory buffer
#
tmem.allocate(self.num_tmem_alloc_cols)
#
# Bar sync for retrieve tensor memory ptr from shared memory
#
tmem.wait_for_alloc()
#
# Retrieving tensor memory ptr and make accumulator tensor
#
acc1_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype)
tCtAcc1_base = cute.make_tensor(acc1_tmem_ptr, tCtAcc_fake.layout)
acc2_tmem_ptr = cute.recast_ptr(
acc1_tmem_ptr + tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base),
dtype=self.acc_dtype,
)
tCtAcc2_base = cute.make_tensor(acc2_tmem_ptr, tCtAcc_fake.layout)
#
# Partition for epilogue
#
epi_tidx = tidx
(
tiled_copy_t2r,
tTR_tAcc1_base,
tTR_tAcc2_base,
tTR_rAcc1,
tTR_rAcc2,
) = self.epilog_tmem_copy_and_partition(
epi_tidx, tCtAcc1_base, tCtAcc2_base, tCgC, epi_tile, use_2cta_instrs
)
tTR_rC = cute.make_rmem_tensor(tTR_rAcc1.shape, self.c_dtype)
tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(
tiled_copy_t2r, tTR_rC, epi_tidx, sC
)
(
tma_atom_c,
bSG_sC,
bSG_gC_partitioned,
) = self.epilog_gmem_copy_and_partition(
epi_tidx, tma_atom_c, tCgC, epi_tile, sC
)
#
# Persistent tile scheduling loop
#
tile_sched = utils.StaticPersistentTileScheduler.create(
tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
)
work_tile = tile_sched.initial_work_tile_info()
acc_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
)
# Threads/warps participating in tma store pipeline
# c_producer_group = pipeline.CooperativeGroup(
# pipeline.Agent.Thread,
# 32 * len(self.epilog_warp_id),
# )
# c_pipeline = pipeline.PipelineTmaStore.create(
# num_stages=self.num_c_stage,
# producer_group=c_producer_group,
# )
while work_tile.is_valid_tile:
# Get tile coord from tile scheduler
cur_tile_coord = work_tile.tile_idx
mma_tile_coord_mnl = (
cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
cur_tile_coord[1],
cur_tile_coord[2],
)
#
# Slice to per mma tile index
#
# ((ATOM_V, REST_V), EPI_M, EPI_N)
bSG_gC = bSG_gC_partitioned[
(
None,
None,
None,
*mma_tile_coord_mnl,
)
]
# Get accumulator stage index
if cutlass.const_expr(self.overlapping_accum):
acc_stage_index = acc_consumer_state.phase
reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False)
else:
acc_stage_index = acc_consumer_state.index
# Set tensor memory buffer for current tile
# (T2R, T2R_M, T2R_N, EPI_M, EPI_M)
tTR_tAcc1 = tTR_tAcc1_base[
(None, None, None, None, None, acc_stage_index)
]
tTR_tAcc2 = tTR_tAcc2_base[
(None, None, None, None, None, acc_stage_index)
]
#
# Wait for accumulator buffer full
#
acc_pipeline.consumer_wait(acc_consumer_state)
tTR_tAcc1 = cute.group_modes(tTR_tAcc1, 3, cute.rank(tTR_tAcc1))
tTR_tAcc2 = cute.group_modes(tTR_tAcc2, 3, cute.rank(tTR_tAcc2))
bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))
#
# Store accumulator to global memory in subtiles
#
subtile_cnt = cute.size(tTR_tAcc1.shape, mode=[3])
num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt
# cute.printf(tTR_tAcc1.shape)
for subtile_idx in cutlass.range(subtile_cnt):
real_subtile_idx = subtile_idx
if cutlass.const_expr(self.overlapping_accum):
if reverse_subtile:
real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n - 1 - subtile_idx
#
# Load accumulator from tensor memory buffer to register
#
tTR_tAcc1_mn = tTR_tAcc1[(None, None, None, real_subtile_idx)]
tTR_tAcc2_mn = tTR_tAcc2[(None, None, None, real_subtile_idx)]
cute.copy(tiled_copy_t2r, tTR_tAcc1_mn, tTR_rAcc1)
cute.copy(tiled_copy_t2r, tTR_tAcc2_mn, tTR_rAcc2)
#
# Async arrive accumulator buffer empty ealier when overlapping_accum is enabled
#
# if cutlass.const_expr(self.overlapping_accum):
# if subtile_idx == self.iter_acc_early_release_in_epilogue:
# # Fence for TMEM load
# cute.arch.fence_view_async_tmem_load()
# with cute.arch.elect_one():
# acc_pipeline.consumer_release(acc_consumer_state)
# acc_consumer_state.advance()
# # silu
acc1_vec = tiled_copy_r2s.retile(tTR_rAcc1).load()
acc2_vec = tiled_copy_r2s.retile(tTR_rAcc2).load()
# acc_vec_test = epilogue_op(acc1_vec) * acc2_vec
acc_vec_test =cute.TensorSSA(
silu_intrinsic(acc1_vec, acc2_vec, cute.size(acc1_vec.shape), appro=self.appro),
acc1_vec.shape,
cutlass.Float32,
)
tRS_rC.store(acc_vec_test.to(self.c_dtype))
#
# Store C to shared memory
#
c_buffer = (num_prev_subtiles + real_subtile_idx) % self.num_c_stage
cute.copy(
tiled_copy_r2s,
tRS_rC,
tRS_sC[(None, None, None, c_buffer)],
)
# Fence and barrier to make sure shared memory store is visible to TMA store
cute.arch.fence_proxy(
cute.arch.ProxyKind.async_shared,
space=cute.arch.SharedSpace.shared_cta,
)
self.epilog_sync_barrier.arrive_and_wait()
#
# TMA store C to global memory
#
if warp_idx == self.epilog_warp_id[0]:
cute.copy(
tma_atom_c,
bSG_sC[(None, c_buffer)],
bSG_gC[(None, real_subtile_idx)],
)
# Fence and barrier to make sure shared memory store is visible to TMA store
# c_pipeline.producer_commit()
# c_pipeline.producer_acquire()
# self.epilog_sync_barrier.arrive_and_wait()
#
# Async arrive accumulator buffer empty
#
if cutlass.const_expr(not self.overlapping_accum):
with cute.arch.elect_one():
acc_pipeline.consumer_release(acc_consumer_state)
acc_consumer_state.advance()
#
# Advance to next tile
#
tile_sched.advance_to_next_work()
work_tile = tile_sched.get_current_work()
#
# Dealloc the tensor memory buffer
#
# tmem.relinquish_alloc_permit()
# self.epilog_sync_barrier.arrive_and_wait()
tmem.free(acc1_tmem_ptr)
#
# Wait for C store complete
#
# c_pipeline.producer_tail()
def mainloop_s2t_copy_and_partition(
self,
sSF: cute.Tensor,
tSF: cute.Tensor,
) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
"""
Make tiledCopy for smem to tmem load for scale factor tensor, then use it to partition smem memory (source) and tensor memory (destination).
:param sSF: The scale factor tensor in smem
:type sSF: cute.Tensor
:param tSF: The scale factor tensor in tmem
:type tSF: cute.Tensor
:return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where:
- tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t)
- tCsSF_compact_s2t: The partitioned scale factor tensor in smem
- tSF_compact_s2t: The partitioned scale factor tensor in tmem
:rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
"""
# (MMA, MMA_MN, MMA_K, STAGE)
tCsSF_compact = cute.filter_zeros(sSF)
# (MMA, MMA_MN, MMA_K)
tCtSF_compact = cute.filter_zeros(tSF)
# Make S2T CopyAtom and tiledCopy
copy_atom_s2t = cute.make_copy_atom(
tcgen05.Cp4x32x128bOp(self.cta_group),
self.sf_dtype,
)
tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact)
thr_copy_s2t = tiled_copy_s2t.get_slice(0)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE)
tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(
tiled_copy_s2t, tCsSF_compact_s2t_
)
# ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K)
tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact)
return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t
def epilog_tmem_copy_and_partition(
self,
tidx: cutlass.Int32,
tAcc1: cute.Tensor,
tAcc2: cute.Tensor,
gC_mnl: cute.Tensor,
epi_tile: cute.Tile,
use_2cta_instrs: Union[cutlass.Boolean, bool],
) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
"""
Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination).
:param tidx: The thread index in epilogue warp groups
:type tidx: cutlass.Int32
:param tAcc: The accumulator tensor to be copied and partitioned
:type tAcc: cute.Tensor
:param gC_mnl: The global tensor C
:type gC_mnl: cute.Tensor
:param epi_tile: The epilogue tiler
:type epi_tile: cute.Tile
:param use_2cta_instrs: Whether use_2cta_instrs is enabled
:type use_2cta_instrs: bool
:return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where:
- tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r)
- tTR_tAcc: The partitioned accumulator tensor
- tTR_rAcc: The accumulated tensor in register used to hold t2r results
:rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
"""
# Make tiledCopy for tensor memory load
copy_atom_t2r = sm100_utils.get_tmem_load_op(
self.cta_tile_shape_mnk,
self.c_layout,
self.c_dtype,
self.acc_dtype,
epi_tile,
use_2cta_instrs,
)
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE)
tAcc1_epi = cute.flat_divide(
tAcc1[((None, None), 0, 0, None)],
epi_tile,
)
tAcc2_epi = cute.flat_divide(
tAcc2[((None, None), 0, 0, None)],
epi_tile,
)
# (EPI_TILE_M, EPI_TILE_N)
tiled_copy_t2r = tcgen05.make_tmem_copy(
copy_atom_t2r, tAcc1_epi[(None, None, 0, 0, 0)]
)
thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
# (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE)
tTR_tAcc1 = thr_copy_t2r.partition_S(tAcc1_epi)
tTR_tAcc2 = thr_copy_t2r.partition_S(tAcc2_epi)
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL)
gC_mnl_epi = cute.flat_divide(
gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
)
# (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL)
tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
# (T2R, T2R_M, T2R_N)
tTR_rAcc1 = cute.make_rmem_tensor(
tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype
)
tTR_rAcc2 = cute.make_rmem_tensor(
tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype
)
return tiled_copy_t2r, tTR_tAcc1, tTR_tAcc2, tTR_rAcc1, tTR_rAcc2
def epilog_smem_copy_and_partition(
self,
tiled_copy_t2r: cute.TiledCopy,
tTR_rC: cute.Tensor,
tidx: cutlass.Int32,
sC: cute.Tensor,
) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]:
"""
Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination).
:param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r)
:type tiled_copy_t2r: cute.TiledCopy
:param tTR_rC: The partitioned accumulator tensor
:type tTR_rC: cute.Tensor
:param tidx: The thread index in epilogue warp groups
:type tidx: cutlass.Int32
:param sC: The shared memory tensor to be copied and partitioned
:type sC: cute.Tensor
:type sepi: cute.Tensor
:return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where:
- tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s)
- tRS_rC: The partitioned tensor C (register source)
- tRS_sC: The partitioned tensor C (smem destination)
:rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]
"""
copy_atom_r2s = sm100_utils.get_smem_store_op(
self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r
)
tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r)
# (R2S, R2S_M, R2S_N, PIPE_D)
thr_copy_r2s = tiled_copy_r2s.get_slice(tidx)
tRS_sC = thr_copy_r2s.partition_D(sC)
# (R2S, R2S_M, R2S_N)
tRS_rC = tiled_copy_r2s.retile(tTR_rC)
return tiled_copy_r2s, tRS_rC, tRS_sC
def epilog_gmem_copy_and_partition(
self,
tidx: cutlass.Int32,
atom: Union[cute.CopyAtom, cute.TiledCopy],
gC_mnl: cute.Tensor,
epi_tile: cute.Tile,
sC: cute.Tensor,
) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]:
"""Make tiledCopy for global memory store, then use it to:
partition shared memory (source) and global memory (destination) for TMA store version.
:param tidx: The thread index in epilogue warp groups
:type tidx: cutlass.Int32
:param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version
:type atom: cute.CopyAtom or cute.TiledCopy
:param gC_mnl: The global tensor C
:type gC_mnl: cute.Tensor
:param epi_tile: The epilogue tiler
:type epi_tile: cute.Tile
:param sC: The shared memory tensor to be copied and partitioned
:type sC: cute.Tensor
:return: A tuple containing (tma_atom_c, bSG_sC, bSG_gC) where:
- tma_atom_c: The TMA copy atom
- bSG_sC: The partitioned shared memory tensor C
- bSG_gC: The partitioned global tensor C
:rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]
"""
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL)
gC_epi = cute.flat_divide(
gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile
)
tma_atom_c = atom
sC_for_tma_partition = cute.group_modes(sC, 0, 2)
gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2)
# ((ATOM_V, REST_V), EPI_M, EPI_N)
# ((ATOM_V, REST_V), EPI_M, EPI_N, RestM, RestN, RestL)
bSG_sC, bSG_gC = cpasync.tma_partition(
tma_atom_c,
0,
cute.make_layout(1),
sC_for_tma_partition,
gC_for_tma_partition,
)
return tma_atom_c, bSG_sC, bSG_gC
@staticmethod
def _compute_stages(
tiled_mma: cute.TiledMma,
mma_tiler_mnk: Tuple[int, int, int],
a_dtype: Type[cutlass.Numeric],
b_dtype: Type[cutlass.Numeric],
epi_tile: cute.Tile,
c_dtype: Type[cutlass.Numeric],
c_layout: utils.LayoutEnum,
sf_dtype: Type[cutlass.Numeric],
sf_vec_size: int,
smem_capacity: int,
occupancy: int,
) -> Tuple[int, int, int]:
"""Computes the number of stages for A/B/C operands based on heuristics.
:param tiled_mma: The tiled MMA object defining the core computation.
:type tiled_mma: cute.TiledMma
:param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler.
:type mma_tiler_mnk: tuple[int, int, int]
:param a_dtype: Data type of operand A.
:type a_dtype: type[cutlass.Numeric]
:param b_dtype: Data type of operand B.
:type b_dtype: type[cutlass.Numeric]
:param epi_tile: The epilogue tile shape.
:type epi_tile: cute.Tile
:param c_dtype: Data type of operand C (output).
:type c_dtype: type[cutlass.Numeric]
:param c_layout: Layout enum of operand C.
:type c_layout: utils.LayoutEnum
:param sf_dtype: Data type of Scale factor.
:type sf_dtype: type[cutlass.Numeric]
:param sf_vec_size: Scale factor vector size.
:type sf_vec_size: int
:param smem_capacity: Total available shared memory capacity in bytes.
:type smem_capacity: int
:param occupancy: Target number of CTAs per SM (occupancy).
:type occupancy: int
:return: A tuple containing the computed number of stages for:
(ACC stages, A/B operand stages, C stages)
:rtype: tuple[int, int, int]
"""
# ACC stages
# num_acc_stage = 1 if mma_tiler_mnk[1] == 256 else 2
num_acc_stage = 1
# Default C stages
num_c_stage = 2
# Calculate smem layout and size for one stage of A, B, SFA, SFB and C
a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(
tiled_mma,
mma_tiler_mnk,
a_dtype,
1, # a tmp 1 stage is provided
)
b_smem_layout_staged_one = sm100_utils.make_smem_layout_b(
tiled_mma,
mma_tiler_mnk,
b_dtype,
1, # a tmp 1 stage is provided
)
sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa(
tiled_mma,
mma_tiler_mnk,
sf_vec_size,
1, # a tmp 1 stage is provided
)
sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb(
tiled_mma,
mma_tiler_mnk,
sf_vec_size,
1, # a tmp 1 stage is provided
)
c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi(
c_dtype,
c_layout,
epi_tile,
1,
)
ab_bytes_per_stage = (
cute.size_in_bytes(a_dtype, a_smem_layout_stage_one)
+ cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) * 2
+ cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one)
+ cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) * 2
)
mbar_helpers_bytes = 1024
c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one)
c_bytes = c_bytes_per_stage * num_c_stage
# Calculate A/B/SFA/SFB stages:
# Start with total smem per CTA (capacity / occupancy)
# Subtract reserved bytes and initial C stages bytes
# Divide remaining by bytes needed per A/B/SFA/SFB stage
num_ab_stage = (
smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes)
) // ab_bytes_per_stage
# Refine epilogue stages:
# Calculate remaining smem after allocating for A/B/SFA/SFB stages and reserved bytes
# Add remaining unused smem to epilogue
num_c_stage += (
smem_capacity
- occupancy * ab_bytes_per_stage * num_ab_stage
- occupancy * (mbar_helpers_bytes + c_bytes)
) // (occupancy * c_bytes_per_stage)
return num_acc_stage, num_ab_stage, num_c_stage
@staticmethod
def _compute_grid(
c: cute.Tensor,
cta_tile_shape_mnk: Tuple[int, int, int],
cluster_shape_mn: Tuple[int, int],
max_active_clusters: cutlass.Constexpr,
) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]:
"""Use persistent tile scheduler to compute the grid size for the output tensor C.
:param c: The output tensor C
:type c: cute.Tensor
:param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile.
:type cta_tile_shape_mnk: tuple[int, int, int]
:param cluster_shape_mn: Shape of each cluster in M, N dimensions.
:type cluster_shape_mn: tuple[int, int]
:param max_active_clusters: Maximum number of active clusters.
:type max_active_clusters: cutlass.Constexpr
:return: A tuple containing:
- tile_sched_params: Parameters for the persistent tile scheduler.
- grid: Grid shape for kernel launch.
:rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]]
"""
c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0))
gc = cute.zipped_divide(c, tiler=c_shape)
num_ctas_mnl = gc[(0, (None, None, None))].shape
cluster_shape_mnl = (*cluster_shape_mn, 1)
tile_sched_params = utils.PersistentTileSchedulerParams(
num_ctas_mnl, cluster_shape_mnl
)
grid = utils.StaticPersistentTileScheduler.get_grid_shape(
tile_sched_params, max_active_clusters
)
return tile_sched_params, grid
# Global cache for compiled kernel
_compiled_kernel_cache = {}
# This function is used to compile the kernel once and cache it and then allow users to
# run the kernel multiple times to get more accurate timing results.
def compile_kernel(problem_size):
"""
Compile the kernel once and cache it.
This should be called before any timing measurements.
Returns:
The compiled kernel function
"""
global _compiled_kernel_cache
if problem_size in _compiled_kernel_cache:
return _compiled_kernel_cache[problem_size]
# Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
a_ptr = make_ptr(
ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16
)
b1_ptr = make_ptr(
ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16
)
b2_ptr = make_ptr(
ab_dtype, 0, cute.AddressSpace.gmem, assumed_align=16
)
c_ptr = make_ptr(
c_dtype, 0, cute.AddressSpace.gmem, assumed_align=16
)
sfa_ptr = make_ptr(
sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32
)
sfb1_ptr = make_ptr(
sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32
)
sfb2_ptr = make_ptr(
sf_dtype, 0, cute.AddressSpace.gmem, assumed_align=32
)
m,n,k,_ = problem_size
gemm = Sm100BlockScaledPersistentDualGemmKernel(
sf_vec_size,
(m,n,k),
)
# Compile the kernel
_compiled_kernel_cache[problem_size] = cute.compile(gemm, a_ptr, b1_ptr, b2_ptr, sfa_ptr, sfb1_ptr, sfb2_ptr, c_ptr, problem_size)
return _compiled_kernel_cache[problem_size]
def custom_kernel(data: input_t) -> output_t:
"""
Execute the block-scaled GEMM kernel.
This is the main entry point called by the evaluation framework.
It converts PyTorch tensors to CuTe tensors, launches the kernel,
and returns the result.
Args:
data: Tuple of (a, b, sfa_ref, sfb_ref, sfa_permuted, sfb_permuted, c) PyTorch tensors
a: [m, k, l] - Input matrix in float4e2m1fn
b: [n, k, l] - Input vector in float4e2m1fn
sfa_ref: [m, k, l] - Scale factors in float8_e4m3fn, used by reference implementation
sfb_ref: [n, k, l] - Scale factors in float8_e4m3fn, used by reference implementation
sfa_permuted: [32, 4, rest_m, 4, rest_k, l] - Scale factors in float8_e4m3fn
sfb_permuted: [32, 4, rest_n, 4, rest_k, l] - Scale factors in float8_e4m3fn
c: [m, n, l] - Output vector in float16
Returns:
Output tensor c with computed results
"""
# a, b, _, _, sfa_permuted, sfb_permuted, c = data
a, b1, b2, _, _, _, sfa_permuted, sfb1_permuted, sfb2_permuted, c = data
# Ensure kernel is compiled (will use cached version if available)
# To avoid the compilation overhead, we compile the kernel once and cache it.
# Get dimensions from MxKxL layout
_, k, _ = a.shape
m, n, l = c.shape
# Torch use e2m1_x2 data type, thus k is halved
k = k * 2
if (m,n,k) in [(256,4096,7168), (512,4096,7168), (256,3072,4096), (512,3072,7168)]: # (512,4096,7168), (256,3072,4096), (512,3072,7168)
compiled_func = compile_kernel((m, n, k, l))
# Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
a_ptr = make_ptr(
ab_dtype, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
)
b1_ptr = make_ptr(
ab_dtype, b1.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
)
b2_ptr = make_ptr(
ab_dtype, b2.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
)
c_ptr = make_ptr(
c_dtype, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=16
)
sfa_ptr = make_ptr(
sf_dtype, sfa_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
sfb1_ptr = make_ptr(
sf_dtype, sfb1_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
sfb2_ptr = make_ptr(
sf_dtype, sfb2_permuted.data_ptr(), cute.AddressSpace.gmem, assumed_align=32
)
# Execute the compiled kernel
compiled_func(a_ptr, b1_ptr, b2_ptr, sfa_ptr, sfb1_ptr, sfb2_ptr, c_ptr)
# compiled_func(a_ptr, b1_ptr, sfa_ptr, sfb1_ptr, c_ptr)
return c
else:
return ref_kernel(data)scrolls · 2352 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 242254.
⋯ 39 unchanged lines# 10.3# 16.5- # problem size: tile_mn | cluster_mn | pf_dist | debug | m512+ # problem size: tile_mn | cluster_mn | pf_distconfig_map = {- (256,4096,7168) : ((256, 64), (2, 1), 0, False, False),- (512,4096,7168) : ((256, 128), (2, 1), 0, True, True),- (256,3072,4096) : ((256, 64), (2, 1), None, False, False),- (512,3072,7168) : ((256, 128), (2, 1), None, False, True),+ (256,4096,7168) : ((256, 64), (2, 1), 0, True),+ (512,4096,7168) : ((256, 128), (2, 1), 0, False),+ (256,3072,4096) : ((256, 64), (2, 1), None, True),+ (512,3072,7168) : ((256, 128), (2, 1), None, True),}ab_dtype = cutlass.Float4E2M1FN⋯ 3 unchanged linesmax_active_clusters = 148@dsl_user_op- def silu_precise_8(src_A, src_B, *, loc=None, ip=None):+ def silu_precise_8(src_A, src_B, appro=False, loc=None, ip=None):inputs = []# 提取 A0-A7for i in range(8):⋯ 17 unchanged linesmul.f32 $0, $16, $0; mul.f32 $1, $17, $1; mul.f32 $2, $18, $2; mul.f32 $3, $19, $3;mul.f32 $4, $20, $4; mul.f32 $5, $21, $5; mul.f32 $6, $22, $6; mul.f32 $7, $23, $7;"""+ if appro:+ asm = r"""+ mul.f32 $0, $8, 0.5; mul.f32 $1, $9, 0.5; mul.f32 $2, $10, 0.5; mul.f32 $3, $11, 0.5;+ mul.f32 $4, $12, 0.5; mul.f32 $5, $13, 0.5; mul.f32 $6, $14, 0.5; mul.f32 $7, $15, 0.5;+ tanh.approx.f32 $0, $0; tanh.approx.f32 $1, $1; tanh.approx.f32 $2, $2; tanh.approx.f32 $3, $3;+ tanh.approx.f32 $4, $4; tanh.approx.f32 $5, $5; tanh.approx.f32 $6, $6; tanh.approx.f32 $7, $7;+ add.f32 $0, $0, 1.0; add.f32 $1, $1, 1.0; add.f32 $2, $2, 1.0; add.f32 $3, $3, 1.0;+ add.f32 $4, $4, 1.0; add.f32 $5, $5, 1.0; add.f32 $6, $6, 1.0; add.f32 $7, $7, 1.0;+ mul.f32 $0, $0, $8; mul.f32 $1, $1, $9; mul.f32 $2, $2, $10; mul.f32 $3, $3, $11;+ mul.f32 $4, $4, $12; mul.f32 $5, $5, $13; mul.f32 $6, $6, $14; mul.f32 $7, $7, $15;+ mul.f32 $0, $0, 0.5; mul.f32 $1, $1, 0.5; mul.f32 $2, $2, 0.5; mul.f32 $3, $3, 0.5;+ mul.f32 $4, $4, 0.5; mul.f32 $5, $5, 0.5; mul.f32 $6, $6, 0.5; mul.f32 $7, $7, 0.5;+ mul.f32 $0, $16, $0; mul.f32 $1, $17, $1; mul.f32 $2, $18, $2; mul.f32 $3, $19, $3;+ mul.f32 $4, $20, $4; mul.f32 $5, $21, $5; mul.f32 $6, $22, $6; mul.f32 $7, $23, $7;+ """++ # asm = r"""+ # mul.f32 $0, $8, 0.5; mul.f32 $1, $9, 0.5; mul.f32 $2, $10, 0.5; mul.f32 $3, $11, 0.5;+ # mul.f32 $4, $12, 0.5; mul.f32 $5, $13, 0.5; mul.f32 $6, $14, 0.5; mul.f32 $7, $15, 0.5;+ # {+ # .reg .f32 %t<8>;+ # tanh.approx.f32 %t0, $0; tanh.approx.f32 %t1, $1; tanh.approx.f32 %t2, $2; tanh.approx.f32 %t3, $3;+ # tanh.approx.f32 %t4, $4; tanh.approx.f32 %t5, $5; tanh.approx.f32 %t6, $6; tanh.approx.f32 %t7, $7;+ # fma.rn.f32 $0, %t0, $0, $0; fma.rn.f32 $1, %t1, $1, $1; fma.rn.f32 $2, %t2, $2, $2; fma.rn.f32 $3, %t3, $3, $3;+ # fma.rn.f32 $4, %t4, $4, $4; fma.rn.f32 $5, %t5, $5, $5; fma.rn.f32 $6, %t6, $6, $6; fma.rn.f32 $7, %t7, $7, $7;+ # }+ # mul.f32 $0, $0, $16; mul.f32 $1, $1, $17; mul.f32 $2, $2, $18; mul.f32 $3, $3, $19;+ # mul.f32 $4, $4, $20; mul.f32 $5, $5, $21; mul.f32 $6, $6, $22; mul.f32 $7, $7, $23;+ # """+ cons = "=&f,=&f,=&f,=&f,=&f,=&f,=&f,=&f,f,f,f,f,f,f,f,f,f,f,f,f,f,f,f,f"- cons = "=f,=f,=f,=f,=f,=f,=f,=f,f,f,f,f,f,f,f,f,f,f,f,f,f,f,f,f"+ # cons = "=f,=f,=f,=f,=f,=f,=f,=f,f,f,f,f,f,f,f,f,f,f,f,f,f,f,f,f"res = llvm.inline_asm(llvm.StructType.get_literal([Float32.mlir_type]*8), inputs, asm, cons, loc=loc, ip=ip)out = []⋯ 2 unchanged linesreturn vector.from_elements(ir.VectorType.get([8], Float32.mlir_type, loc=loc), out, loc=loc, ip=ip)@dsl_user_op- def silu_intrinsic(vec_A, vec_B, length, *, loc=None, ip=None):+ def silu_intrinsic(vec_A, vec_B, length, appro=False, loc=None, ip=None):src_pos = 0vec_f32x8_type = ir.VectorType.get([8], Float32.mlir_type, loc=loc)vec_dst_type = ir.VectorType.get([length], Float32.mlir_type, loc=loc)⋯ 8 unchanged lines)vec_dst = vector.insert_strided_slice(- silu_precise_8(vec_f32x8_A, vec_f32x8_B),+ silu_precise_8(vec_f32x8_A, vec_f32x8_B, appro),vec_dst,[src_pos],[1],⋯ 145 unchanged lines"""mma_tiler_mn = config_map[problem_size][0]cluster_shape_mn = config_map[problem_size][1]- self.debug = config_map[problem_size][3]- self.m512 = config_map[problem_size][4]+ self.appro = config_map[problem_size][3]+self.acc_dtype = cutlass.Float32self.sf_vec_size = sf_vec_sizeself.use_2cta_instrs = mma_tiler_mn[0] == 256⋯ 212 unchanged linessfb2_ptr: cute.Pointer,c_ptr: cute.Pointer,problem_size: cutlass.Constexpr,- epilogue_op: cutlass.Constexpr = lambda x: x,+ # epilogue_op: cutlass.Constexpr = lambda x: x,+ epilogue_op: cutlass.Constexpr = lambda x: 0.5*x * (1.0 + cute.math.tanh(0.5*x, fastmath=True)),):m, n, k, l = problem_size⋯ 672 unchanged lines## Persistent tile scheduling loop#- # tile_sched = utils.StaticPersistentTileScheduler.create(- # tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()- # )- # work_tile = tile_sched.initial_work_tile_info()+ tile_sched = utils.StaticPersistentTileScheduler.create(+ tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()+ )+ work_tile = tile_sched.initial_work_tile_info()ab_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_ab_stage)- # while work_tile.is_valid_tile:- # Get tile coord from tile scheduler- if cutlass.const_expr(self.m512):- cur_tile_coord = ((bidx + bidz * 2) % 4 , bidz // 2, 0)- else:- cur_tile_coord = (bidx, bidz, 0)-- mma_tile_coord_mnl = (- cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),- cur_tile_coord[1],- cur_tile_coord[2],- )+ while work_tile.is_valid_tile:+ # Get tile coord from tile scheduler+ cur_tile_coord = work_tile.tile_idx+ mma_tile_coord_mnl = (+ cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),+ cur_tile_coord[1],+ cur_tile_coord[2],+ )- #- # Slice to per mma tile index- #- # ((atom_v, rest_v), RestK)- tAgA_slice = tAgA[- (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])- ]- # ((atom_v, rest_v), RestK)- tBgB1_slice = tBgB1[- (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])- ]- tBgB2_slice = tBgB2[- (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])- ]+ #+ # Slice to per mma tile index+ #+ # ((atom_v, rest_v), RestK)+ tAgA_slice = tAgA[+ (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])+ ]+ # ((atom_v, rest_v), RestK)+ tBgB1_slice = tBgB1[+ (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])+ ]+ tBgB2_slice = tBgB2[+ (None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])+ ]- # ((atom_v, rest_v), RestK)- tAgSFA_slice = tAgSFA[- (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])- ]+ # ((atom_v, rest_v), RestK)+ tAgSFA_slice = tAgSFA[+ (None, mma_tile_coord_mnl[0], None, mma_tile_coord_mnl[2])+ ]- slice_n = mma_tile_coord_mnl[1]- if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):- slice_n = mma_tile_coord_mnl[1] // 2- # ((atom_v, rest_v), RestK)- tBgSFB1_slice = tBgSFB1[- (None, slice_n, None, mma_tile_coord_mnl[2])- ]- tBgSFB2_slice = tBgSFB2[- (None, slice_n, None, mma_tile_coord_mnl[2])- ]+ slice_n = mma_tile_coord_mnl[1]+ if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):+ slice_n = mma_tile_coord_mnl[1] // 2+ # ((atom_v, rest_v), RestK)+ tBgSFB1_slice = tBgSFB1[+ (None, slice_n, None, mma_tile_coord_mnl[2])+ ]+ tBgSFB2_slice = tBgSFB2[+ (None, slice_n, None, mma_tile_coord_mnl[2])+ ]- #- # Prefetch: Initial batch of prefetches to prime the pipeline- #- if self.prefetch_enabled:- for pf_k_tile in cutlass.range(- 0, min(self.prefetch_dist, k_tile_cnt), unroll=1- ):- cute.prefetch(- tma_atom_a,- tAgA_slice[(None, pf_k_tile)],- )- cute.prefetch(- tma_atom_b1,- tBgB1_slice[(None, pf_k_tile)],- )- cute.prefetch(- tma_atom_b2,- tBgB2_slice[(None, pf_k_tile)],- )- cute.prefetch(- tma_atom_sfa,- tAgSFA_slice[(None, pf_k_tile)],- )- cute.prefetch(- tma_atom_sfb1,- tBgSFB1_slice[(None, pf_k_tile)],- )- cute.prefetch(- tma_atom_sfb2,- tBgSFB2_slice[(None, pf_k_tile)],- )-- # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt- ab_producer_state.reset_count()- peek_ab_empty_status = cutlass.Boolean(1)- if ab_producer_state.count < k_tile_cnt:- peek_ab_empty_status = ab_pipeline.producer_try_acquire(- ab_producer_state- )- #- # Tma load loop- #- for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):- # Conditionally wait for AB buffer empty- ab_pipeline.producer_acquire(- ab_producer_state, peek_ab_empty_status- )-- # TMA load A/B/SFA/SFB- cute.copy(- tma_atom_a,- tAgA_slice[(None, ab_producer_state.count)],- tAsA[(None, ab_producer_state.index)],- tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),- mcast_mask=a_full_mcast_mask,- )- cute.copy(- tma_atom_b1,- tBgB1_slice[(None, ab_producer_state.count)],- tBsB1[(None, ab_producer_state.index)],- tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),- mcast_mask=b_full_mcast_mask,- )- cute.copy(- tma_atom_b2,- tBgB2_slice[(None, ab_producer_state.count)],- tBsB2[(None, ab_producer_state.index)],- tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),- mcast_mask=b_full_mcast_mask,- )- cute.copy(- tma_atom_sfa,- tAgSFA_slice[(None, ab_producer_state.count)],- tAsSFA[(None, ab_producer_state.index)],- tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),- mcast_mask=sfa_full_mcast_mask,- )- cute.copy(- tma_atom_sfb1,- tBgSFB1_slice[(None, ab_producer_state.count)],- tBsSFB1[(None, ab_producer_state.index)],- tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),- mcast_mask=sfb_full_mcast_mask,- )- cute.copy(- tma_atom_sfb2,- tBgSFB2_slice[(None, ab_producer_state.count)],- tBsSFB2[(None, ab_producer_state.index)],- tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),- mcast_mask=sfb_full_mcast_mask,- )-- # Prefetch: Rolling prefetch for next tiles+ #+ # Prefetch: Initial batch of prefetches to prime the pipeline+ #if self.prefetch_enabled:- if k_tile < k_tile_cnt - self.prefetch_dist:- future_k_tile = ab_producer_state.count + self.prefetch_dist+ for pf_k_tile in cutlass.range(+ 0, min(self.prefetch_dist, k_tile_cnt), unroll=1+ ):cute.prefetch(tma_atom_a,- tAgA_slice[(None, future_k_tile)],+ tAgA_slice[(None, pf_k_tile)],)cute.prefetch(tma_atom_b1,- tBgB1_slice[(None, future_k_tile)],+ tBgB1_slice[(None, pf_k_tile)],)cute.prefetch(tma_atom_b2,- tBgB2_slice[(None, future_k_tile)],+ tBgB2_slice[(None, pf_k_tile)],)cute.prefetch(tma_atom_sfa,- tAgSFA_slice[(None, future_k_tile)],+ tAgSFA_slice[(None, pf_k_tile)],)cute.prefetch(tma_atom_sfb1,- tBgSFB1_slice[(None, future_k_tile)],+ tBgSFB1_slice[(None, pf_k_tile)],)cute.prefetch(tma_atom_sfb2,- tBgSFB2_slice[(None, future_k_tile)],+ tBgSFB2_slice[(None, pf_k_tile)],)- # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1- ab_producer_state.advance()+ # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt+ ab_producer_state.reset_count()peek_ab_empty_status = cutlass.Boolean(1)if ab_producer_state.count < k_tile_cnt:peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state)+ #+ # Tma load loop+ #+ for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1):+ # Conditionally wait for AB buffer empty+ ab_pipeline.producer_acquire(+ ab_producer_state, peek_ab_empty_status+ )- #- # Advance to next tile- #- # tile_sched.advance_to_next_work()- # work_tile = tile_sched.get_current_work()+ # TMA load A/B/SFA/SFB+ cute.copy(+ tma_atom_a,+ tAgA_slice[(None, ab_producer_state.count)],+ tAsA[(None, ab_producer_state.index)],+ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),+ mcast_mask=a_full_mcast_mask,+ )+ cute.copy(+ tma_atom_b1,+ tBgB1_slice[(None, ab_producer_state.count)],+ tBsB1[(None, ab_producer_state.index)],+ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),+ mcast_mask=b_full_mcast_mask,+ )+ cute.copy(+ tma_atom_b2,+ tBgB2_slice[(None, ab_producer_state.count)],+ tBsB2[(None, ab_producer_state.index)],+ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),+ mcast_mask=b_full_mcast_mask,+ )+ cute.copy(+ tma_atom_sfa,+ tAgSFA_slice[(None, ab_producer_state.count)],+ tAsSFA[(None, ab_producer_state.index)],+ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),+ mcast_mask=sfa_full_mcast_mask,+ )+ cute.copy(+ tma_atom_sfb1,+ tBgSFB1_slice[(None, ab_producer_state.count)],+ tBsSFB1[(None, ab_producer_state.index)],+ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),+ mcast_mask=sfb_full_mcast_mask,+ )+ cute.copy(+ tma_atom_sfb2,+ tBgSFB2_slice[(None, ab_producer_state.count)],+ tBsSFB2[(None, ab_producer_state.index)],+ tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state),+ mcast_mask=sfb_full_mcast_mask,+ )+ # Prefetch: Rolling prefetch for next tiles+ if self.prefetch_enabled:+ if k_tile < k_tile_cnt - self.prefetch_dist:+ future_k_tile = ab_producer_state.count + self.prefetch_dist+ cute.prefetch(+ tma_atom_a,+ tAgA_slice[(None, future_k_tile)],+ )+ cute.prefetch(+ tma_atom_b1,+ tBgB1_slice[(None, future_k_tile)],+ )+ cute.prefetch(+ tma_atom_b2,+ tBgB2_slice[(None, future_k_tile)],+ )+ cute.prefetch(+ tma_atom_sfa,+ tAgSFA_slice[(None, future_k_tile)],+ )+ cute.prefetch(+ tma_atom_sfb1,+ tBgSFB1_slice[(None, future_k_tile)],+ )+ cute.prefetch(+ tma_atom_sfb2,+ tBgSFB2_slice[(None, future_k_tile)],+ )++ # Peek (try_wait) AB buffer empty for k_tile = prefetch_k_tile_cnt + k_tile + 1+ ab_producer_state.advance()+ peek_ab_empty_status = cutlass.Boolean(1)+ if ab_producer_state.count < k_tile_cnt:+ peek_ab_empty_status = ab_pipeline.producer_try_acquire(+ ab_producer_state+ )++ #+ # Advance to next tile+ #+ tile_sched.advance_to_next_work()+ work_tile = tile_sched.get_current_work()+## Wait A/B buffer empty#⋯ 82 unchanged lines## Persistent tile scheduling loop#- # tile_sched = utils.StaticPersistentTileScheduler.create(- # tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()- # )- # work_tile = tile_sched.initial_work_tile_info()+ tile_sched = utils.StaticPersistentTileScheduler.create(+ tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()+ )+ work_tile = tile_sched.initial_work_tile_info()ab_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_ab_stage⋯ 2 unchanged linespipeline.PipelineUserType.Producer, self.num_acc_stage)- # while work_tile.is_valid_tile:- # Get tile coord from tile scheduler- if cutlass.const_expr(self.m512):- cur_tile_coord = ((bidx + bidz * 2) % 4 , bidz // 2, 0)- else:- cur_tile_coord = (bidx, bidz, 0)-- mma_tile_coord_mnl = (- cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),- cur_tile_coord[1],- cur_tile_coord[2],- )-- # Get accumulator stage index- # if cutlass.const_expr(self.overlapping_accum):- # acc_stage_index = acc_producer_state.phase ^ 1- # else:- # acc_stage_index = acc_producer_state.index-- # Set tensor memory buffer for current tile- # (MMA, MMA_M, MMA_N)- tCtAcc1 = tCtAcc1_base[(None, None, None, 0)]- tCtAcc2 = tCtAcc2_base[(None, None, None, 0)]-- # Peek (try_wait) AB buffer full for k_tile = 0- ab_consumer_state.reset_count()- peek_ab_full_status = cutlass.Boolean(1)- if ab_consumer_state.count < k_tile_cnt and is_leader_cta:- peek_ab_full_status = ab_pipeline.consumer_try_wait(- ab_consumer_state+ while work_tile.is_valid_tile:+ # Get tile coord from tile scheduler+ cur_tile_coord = work_tile.tile_idx+ mma_tile_coord_mnl = (+ cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),+ cur_tile_coord[1],+ cur_tile_coord[2],)- #- # Wait for accumulator buffer empty- #- # if is_leader_cta:- # acc_pipeline.producer_acquire(acc_producer_state)+ # Get accumulator stage index+ if cutlass.const_expr(self.overlapping_accum):+ acc_stage_index = acc_producer_state.phase ^ 1+ else:+ acc_stage_index = acc_producer_state.index- tCtSFB1_mma = tCtSFB1- tCtSFB2_mma = tCtSFB2- if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):- # If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words (ignores first 64 columns of SFB)- offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0)- shifted_ptr1 = cute.recast_ptr(- acc1_tmem_ptr- + tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2- + tcgen05.find_tmem_tensor_col_offset(tCtSFA)- + offset,- dtype=self.sf_dtype,- )- shifted_ptr2 = cute.recast_ptr(- acc1_tmem_ptr- + tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2- + tcgen05.find_tmem_tensor_col_offset(tCtSFA)- + tcgen05.find_tmem_tensor_col_offset(tCtSFB1)- + offset,- dtype=self.sf_dtype,- )- tCtSFB1_mma = cute.make_tensor(shifted_ptr1, tCtSFB_layout)- tCtSFB2_mma = cute.make_tensor(shifted_ptr2, tCtSFB_layout)- elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):- # Move in increments of 64 columns of SFB- offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)- shifted_ptr1 = cute.recast_ptr(- acc1_tmem_ptr- + tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2- + tcgen05.find_tmem_tensor_col_offset(tCtSFA)- + offset,- dtype=self.sf_dtype,- )- shifted_ptr2 = cute.recast_ptr(- acc1_tmem_ptr- + tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2- + tcgen05.find_tmem_tensor_col_offset(tCtSFA)- + tcgen05.find_tmem_tensor_col_offset(tCtSFB1)- + offset,- dtype=self.sf_dtype,- )- tCtSFB1_mma = cute.make_tensor(shifted_ptr1, tCtSFB_layout)- tCtSFB2_mma = cute.make_tensor(shifted_ptr2, tCtSFB_layout)+ # Set tensor memory buffer for current tile+ # (MMA, MMA_M, MMA_N)+ tCtAcc1 = tCtAcc1_base[(None, None, None, acc_producer_state.index)]+ tCtAcc2 = tCtAcc2_base[(None, None, None, acc_producer_state.index)]- #- # Reset the ACCUMULATE field for each tile- #- tiled_mma.set(tcgen05.Field.ACCUMULATE, False)+ # Peek (try_wait) AB buffer full for k_tile = 0+ ab_consumer_state.reset_count()+ peek_ab_full_status = cutlass.Boolean(1)+ if ab_consumer_state.count < k_tile_cnt and is_leader_cta:+ peek_ab_full_status = ab_pipeline.consumer_try_wait(+ ab_consumer_state+ )- #- # Mma mainloop- #- for k_tile in range(k_tile_cnt):+ #+ # Wait for accumulator buffer empty+ #if is_leader_cta:- # Conditionally wait for AB buffer full- ab_pipeline.consumer_wait(- ab_consumer_state, peek_ab_full_status- )+ acc_pipeline.producer_acquire(acc_producer_state)- # Copy SFA/SFB from smem to tmem- s2t_stage_coord = (- None,- None,- None,- None,- ab_consumer_state.index,+ tCtSFB1_mma = tCtSFB1+ tCtSFB2_mma = tCtSFB2+ if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192):+ # If this is an ODD tile, shift the TMEM start address for cta_tile_shape_n=192 case by two words (ignores first 64 columns of SFB)+ offset = cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0)+ shifted_ptr1 = cute.recast_ptr(+ acc1_tmem_ptr+ + tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2+ + tcgen05.find_tmem_tensor_col_offset(tCtSFA)+ + offset,+ dtype=self.sf_dtype,)- tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]- tCsSFB1_compact_s2t_staged = tCsSFB1_compact_s2t[s2t_stage_coord]- tCsSFB2_compact_s2t_staged = tCsSFB2_compact_s2t[s2t_stage_coord]- cute.copy(- tiled_copy_s2t_sfa,- tCsSFA_compact_s2t_staged,- tCtSFA_compact_s2t,+ shifted_ptr2 = cute.recast_ptr(+ acc1_tmem_ptr+ + tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2+ + tcgen05.find_tmem_tensor_col_offset(tCtSFA)+ + tcgen05.find_tmem_tensor_col_offset(tCtSFB1)+ + offset,+ dtype=self.sf_dtype,)- cute.copy(- tiled_copy_s2t_sfb1,- tCsSFB1_compact_s2t_staged,- tCtSFB1_compact_s2t,+ tCtSFB1_mma = cute.make_tensor(shifted_ptr1, tCtSFB_layout)+ tCtSFB2_mma = cute.make_tensor(shifted_ptr2, tCtSFB_layout)+ elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64):+ # Move in increments of 64 columns of SFB+ offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2)+ shifted_ptr1 = cute.recast_ptr(+ acc1_tmem_ptr+ + tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2+ + tcgen05.find_tmem_tensor_col_offset(tCtSFA)+ + offset,+ dtype=self.sf_dtype,)- cute.copy(- tiled_copy_s2t_sfb2,- tCsSFB2_compact_s2t_staged,- tCtSFB2_compact_s2t,+ shifted_ptr2 = cute.recast_ptr(+ acc1_tmem_ptr+ + tcgen05.find_tmem_tensor_col_offset(tCtAcc1_base) * 2+ + tcgen05.find_tmem_tensor_col_offset(tCtSFA)+ + tcgen05.find_tmem_tensor_col_offset(tCtSFB1)+ + offset,+ dtype=self.sf_dtype,)+ tCtSFB1_mma = cute.make_tensor(shifted_ptr1, tCtSFB_layout)+ tCtSFB2_mma = cute.make_tensor(shifted_ptr2, tCtSFB_layout)- # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB- num_kblocks = cute.size(tCrA, mode=[2])- for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):- kblock_coord = (+ #+ # Reset the ACCUMULATE field for each tile+ #+ tiled_mma.set(tcgen05.Field.ACCUMULATE, False)++ #+ # Mma mainloop+ #+ for k_tile in range(k_tile_cnt):+ if is_leader_cta:+ # Conditionally wait for AB buffer full+ ab_pipeline.consumer_wait(+ ab_consumer_state, peek_ab_full_status+ )++ # Copy SFA/SFB from smem to tmem+ s2t_stage_coord = (None,None,- kblock_idx,+ None,+ None,ab_consumer_state.index,)-- # Set SFA/SFB tensor to tiled_mma- sf_kblock_coord = (None, None, kblock_idx)- tiled_mma.set(- tcgen05.Field.SFA,- tCtSFA[sf_kblock_coord].iterator,+ tCsSFA_compact_s2t_staged = tCsSFA_compact_s2t[s2t_stage_coord]+ tCsSFB1_compact_s2t_staged = tCsSFB1_compact_s2t[s2t_stage_coord]+ tCsSFB2_compact_s2t_staged = tCsSFB2_compact_s2t[s2t_stage_coord]+ cute.copy(+ tiled_copy_s2t_sfa,+ tCsSFA_compact_s2t_staged,+ tCtSFA_compact_s2t,)- tiled_mma.set(- tcgen05.Field.SFB,- tCtSFB1_mma[sf_kblock_coord].iterator,+ cute.copy(+ tiled_copy_s2t_sfb1,+ tCsSFB1_compact_s2t_staged,+ tCtSFB1_compact_s2t,)-- cute.gemm(- tiled_mma,- tCtAcc1,- tCrA[kblock_coord],- tCrB1[kblock_coord],- tCtAcc1,+ cute.copy(+ tiled_copy_s2t_sfb2,+ tCsSFB2_compact_s2t_staged,+ tCtSFB2_compact_s2t,)- tiled_mma.set(- tcgen05.Field.SFB,- tCtSFB2_mma[sf_kblock_coord].iterator,- )+ # tCtAcc += tCrA * tCrSFA * tCrB * tCrSFB+ num_kblocks = cute.size(tCrA, mode=[2])+ for kblock_idx in cutlass.range(num_kblocks, unroll_full=True):+ kblock_coord = (+ None,+ None,+ kblock_idx,+ ab_consumer_state.index,+ )- cute.gemm(- tiled_mma,- tCtAcc2,- tCrA[kblock_coord],- tCrB2[kblock_coord],- tCtAcc2,- )+ # Set SFA/SFB tensor to tiled_mma+ sf_kblock_coord = (None, None, kblock_idx)+ tiled_mma.set(+ tcgen05.Field.SFA,+ tCtSFA[sf_kblock_coord].iterator,+ )+ tiled_mma.set(+ tcgen05.Field.SFB,+ tCtSFB1_mma[sf_kblock_coord].iterator,+ )- # Enable accumulate on tCtAcc after first kblock- tiled_mma.set(tcgen05.Field.ACCUMULATE, True)+ cute.gemm(+ tiled_mma,+ tCtAcc1,+ tCrA[kblock_coord],+ tCrB1[kblock_coord],+ tCtAcc1,+ )- # Async arrive AB buffer empty- ab_pipeline.consumer_release(ab_consumer_state)+ tiled_mma.set(+ tcgen05.Field.SFB,+ tCtSFB2_mma[sf_kblock_coord].iterator,+ )- # Peek (try_wait) AB buffer full for k_tile = k_tile + 1- ab_consumer_state.advance()- peek_ab_full_status = cutlass.Boolean(1)- if ab_consumer_state.count < k_tile_cnt:- if is_leader_cta:- peek_ab_full_status = ab_pipeline.consumer_try_wait(- ab_consumer_state- )+ cute.gemm(+ tiled_mma,+ tCtAcc2,+ tCrA[kblock_coord],+ tCrB2[kblock_coord],+ tCtAcc2,+ )- #- # Async arrive accumulator buffer full- #- if is_leader_cta:- acc_pipeline.producer_commit(acc_producer_state)- # acc_producer_state.advance()+ # Enable accumulate on tCtAcc after first kblock+ tiled_mma.set(tcgen05.Field.ACCUMULATE, True)- #- # Advance to next tile- #- # tile_sched.advance_to_next_work()- # work_tile = tile_sched.get_current_work()+ # Async arrive AB buffer empty+ ab_pipeline.consumer_release(ab_consumer_state)+ # Peek (try_wait) AB buffer full for k_tile = k_tile + 1+ ab_consumer_state.advance()+ peek_ab_full_status = cutlass.Boolean(1)+ if ab_consumer_state.count < k_tile_cnt:+ if is_leader_cta:+ peek_ab_full_status = ab_pipeline.consumer_try_wait(+ ab_consumer_state+ )++ #+ # Async arrive accumulator buffer full+ #+ if is_leader_cta:+ acc_pipeline.producer_commit(acc_producer_state)+ acc_producer_state.advance()++ #+ # Advance to next tile+ #+ tile_sched.advance_to_next_work()+ work_tile = tile_sched.get_current_work()+## Wait for accumulator buffer empty#- # acc_pipeline.producer_tail(acc_producer_state)+ acc_pipeline.producer_tail(acc_producer_state)## Specialized epilogue warps#⋯ 48 unchanged lines## Persistent tile scheduling loop#- # tile_sched = utils.StaticPersistentTileScheduler.create(- # tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()- # )- # work_tile = tile_sched.initial_work_tile_info()+ tile_sched = utils.StaticPersistentTileScheduler.create(+ tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()+ )+ work_tile = tile_sched.initial_work_tile_info()acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage⋯ 9 unchanged lines# producer_group=c_producer_group,# )- # while work_tile.is_valid_tile:- # Get tile coord from tile scheduler- if cutlass.const_expr(self.m512):- cur_tile_coord = ((bidx + bidz * 2) % 4 , bidz // 2, 0)- else:- cur_tile_coord = (bidx, bidz, 0)-- mma_tile_coord_mnl = (- cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),- cur_tile_coord[1],- cur_tile_coord[2],- )-- #- # Slice to per mma tile index- #- # ((ATOM_V, REST_V), EPI_M, EPI_N)- bSG_gC = bSG_gC_partitioned[- (- None,- None,- None,- *mma_tile_coord_mnl,+ while work_tile.is_valid_tile:+ # Get tile coord from tile scheduler+ cur_tile_coord = work_tile.tile_idx+ mma_tile_coord_mnl = (+ cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),+ cur_tile_coord[1],+ cur_tile_coord[2],)- ]- # Get accumulator stage index- if cutlass.const_expr(self.overlapping_accum):- acc_stage_index = acc_consumer_state.phase- reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False)- else:- acc_stage_index = acc_consumer_state.index-- # Set tensor memory buffer for current tile- # (T2R, T2R_M, T2R_N, EPI_M, EPI_M)- tTR_tAcc1 = tTR_tAcc1_base[- (None, None, None, None, None, acc_stage_index)- ]- tTR_tAcc2 = tTR_tAcc2_base[- (None, None, None, None, None, acc_stage_index)- ]- #- # Wait for accumulator buffer full- #- acc_pipeline.consumer_wait(acc_consumer_state)-- tTR_tAcc1 = cute.group_modes(tTR_tAcc1, 3, cute.rank(tTR_tAcc1))- tTR_tAcc2 = cute.group_modes(tTR_tAcc2, 3, cute.rank(tTR_tAcc2))- bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))-- #- # Store accumulator to global memory in subtiles- #- subtile_cnt = cute.size(tTR_tAcc1.shape, mode=[3])- # num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt- num_prev_subtiles = 0- for subtile_idx in cutlass.range(subtile_cnt):- real_subtile_idx = subtile_idx- if cutlass.const_expr(self.overlapping_accum):- if reverse_subtile:- real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n - 1 - subtile_idx#- # Load accumulator from tensor memory buffer to register+ # Slice to per mma tile index#- tTR_tAcc1_mn = tTR_tAcc1[(None, None, None, real_subtile_idx)]- tTR_tAcc2_mn = tTR_tAcc2[(None, None, None, real_subtile_idx)]- cute.copy(tiled_copy_t2r, tTR_tAcc1_mn, tTR_rAcc1)- cute.copy(tiled_copy_t2r, tTR_tAcc2_mn, tTR_rAcc2)+ # ((ATOM_V, REST_V), EPI_M, EPI_N)+ bSG_gC = bSG_gC_partitioned[+ (+ None,+ None,+ None,+ *mma_tile_coord_mnl,+ )+ ]- #- # Async arrive accumulator buffer empty ealier when overlapping_accum is enabled- #+ # Get accumulator stage indexif cutlass.const_expr(self.overlapping_accum):- if subtile_idx == self.iter_acc_early_release_in_epilogue:- # Fence for TMEM load- cute.arch.fence_view_async_tmem_load()- with cute.arch.elect_one():- acc_pipeline.consumer_release(acc_consumer_state)- acc_consumer_state.advance()+ acc_stage_index = acc_consumer_state.phase+ reverse_subtile = cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False)+ else:+ acc_stage_index = acc_consumer_state.index- # silu- acc1_vec = tiled_copy_r2s.retile(tTR_rAcc1).load()- acc2_vec = tiled_copy_r2s.retile(tTR_rAcc2).load()- acc_vec_test =cute.TensorSSA(- silu_intrinsic(acc1_vec, acc2_vec, cute.size(acc1_vec.shape)),- acc1_vec.shape,- cutlass.Float32,- )- tRS_rC.store(acc_vec_test.to(self.c_dtype))-+ # Set tensor memory buffer for current tile+ # (T2R, T2R_M, T2R_N, EPI_M, EPI_M)+ tTR_tAcc1 = tTR_tAcc1_base[+ (None, None, None, None, None, acc_stage_index)+ ]+ tTR_tAcc2 = tTR_tAcc2_base[+ (None, None, None, None, None, acc_stage_index)+ ]#- # Store C to shared memory+ # Wait for accumulator buffer full#- c_buffer = (num_prev_subtiles + real_subtile_idx) % self.num_c_stage- cute.copy(- tiled_copy_r2s,- tRS_rC,- tRS_sC[(None, None, None, c_buffer)],- )- # Fence and barrier to make sure shared memory store is visible to TMA store- cute.arch.fence_proxy(- cute.arch.ProxyKind.async_shared,- space=cute.arch.SharedSpace.shared_cta,- )- self.epilog_sync_barrier.arrive_and_wait()+ acc_pipeline.consumer_wait(acc_consumer_state)+ tTR_tAcc1 = cute.group_modes(tTR_tAcc1, 3, cute.rank(tTR_tAcc1))+ tTR_tAcc2 = cute.group_modes(tTR_tAcc2, 3, cute.rank(tTR_tAcc2))+ bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC))+#- # TMA store C to global memory+ # Store accumulator to global memory in subtiles#- if warp_idx == self.epilog_warp_id[0]:+ subtile_cnt = cute.size(tTR_tAcc1.shape, mode=[3])+ num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt++ # cute.printf(tTR_tAcc1.shape)+ for subtile_idx in cutlass.range(subtile_cnt):+ real_subtile_idx = subtile_idx+ if cutlass.const_expr(self.overlapping_accum):+ if reverse_subtile:+ real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n - 1 - subtile_idx+ #+ # Load accumulator from tensor memory buffer to register+ #+ tTR_tAcc1_mn = tTR_tAcc1[(None, None, None, real_subtile_idx)]+ tTR_tAcc2_mn = tTR_tAcc2[(None, None, None, real_subtile_idx)]+ cute.copy(tiled_copy_t2r, tTR_tAcc1_mn, tTR_rAcc1)+ cute.copy(tiled_copy_t2r, tTR_tAcc2_mn, tTR_rAcc2)++ #+ # Async arrive accumulator buffer empty ealier when overlapping_accum is enabled+ #+ # if cutlass.const_expr(self.overlapping_accum):+ # if subtile_idx == self.iter_acc_early_release_in_epilogue:+ # # Fence for TMEM load+ # cute.arch.fence_view_async_tmem_load()+ # with cute.arch.elect_one():+ # acc_pipeline.consumer_release(acc_consumer_state)+ # acc_consumer_state.advance()++ # # silu+ acc1_vec = tiled_copy_r2s.retile(tTR_rAcc1).load()+ acc2_vec = tiled_copy_r2s.retile(tTR_rAcc2).load()+ # acc_vec_test = epilogue_op(acc1_vec) * acc2_vec+ acc_vec_test =cute.TensorSSA(+ silu_intrinsic(acc1_vec, acc2_vec, cute.size(acc1_vec.shape), appro=self.appro),+ acc1_vec.shape,+ cutlass.Float32,+ )+ tRS_rC.store(acc_vec_test.to(self.c_dtype))++ #+ # Store C to shared memory+ #+ c_buffer = (num_prev_subtiles + real_subtile_idx) % self.num_c_stagecute.copy(- tma_atom_c,- bSG_sC[(None, c_buffer)],- bSG_gC[(None, real_subtile_idx)],+ tiled_copy_r2s,+ tRS_rC,+ tRS_sC[(None, None, None, c_buffer)],)# Fence and barrier to make sure shared memory store is visible to TMA store- # c_pipeline.producer_commit()- # c_pipeline.producer_acquire()- # self.epilog_sync_barrier.arrive_and_wait()+ cute.arch.fence_proxy(+ cute.arch.ProxyKind.async_shared,+ space=cute.arch.SharedSpace.shared_cta,+ )+ self.epilog_sync_barrier.arrive_and_wait()- #- # Async arrive accumulator buffer empty- #- # if cutlass.const_expr(not self.overlapping_accum):- # with cute.arch.elect_one():- # acc_pipeline.consumer_release(acc_consumer_state)- # acc_consumer_state.advance()+ #+ # TMA store C to global memory+ #+ if warp_idx == self.epilog_warp_id[0]:+ cute.copy(+ tma_atom_c,+ bSG_sC[(None, c_buffer)],+ bSG_gC[(None, real_subtile_idx)],+ )+ # Fence and barrier to make sure shared memory store is visible to TMA store+ # c_pipeline.producer_commit()+ # c_pipeline.producer_acquire()+ # self.epilog_sync_barrier.arrive_and_wait()- #- # Advance to next tile- #- # tile_sched.advance_to_next_work()- # work_tile = tile_sched.get_current_work()+ #+ # Async arrive accumulator buffer empty+ #+ if cutlass.const_expr(not self.overlapping_accum):+ with cute.arch.elect_one():+ acc_pipeline.consumer_release(acc_consumer_state)+ acc_consumer_state.advance()+ #+ # Advance to next tile+ #+ tile_sched.advance_to_next_work()+ work_tile = tile_sched.get_current_work()+## Dealloc the tensor memory buffer#⋯ 442 unchanged lines# Torch use e2m1_x2 data type, thus k is halvedk = k * 2- # (256,4096,7168), (512,4096,7168), (256,3072,4096), (512,3072,7168)- if (m,n,k) in [(256,4096,7168), (512,4096,7168), (256,3072,4096), (512,3072,7168)]:++ if (m,n,k) in [(256,4096,7168), (512,4096,7168), (256,3072,4096), (512,3072,7168)]: # (512,4096,7168), (256,3072,4096), (512,3072,7168)compiled_func = compile_kernel((m, n, k, l))# Create CuTe pointers for A/B/C/SFA/SFB via torch tensor data pointer
scrolls · 1113 diff lines total
Best evidence level for this revision: reported
JSON