3778 lines
151 KiB
Python
3778 lines
151 KiB
Python
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
# Redistribution and use in source and binary forms, with or without
|
|
# modification, are permitted provided that the following conditions are met:
|
|
|
|
# 1. Redistributions of source code must retain the above copyright notice, this
|
|
# list of conditions and the following disclaimer.
|
|
|
|
# 2. Redistributions in binary form must reproduce the above copyright notice,
|
|
# this list of conditions and the following disclaimer in the documentation
|
|
# and/or other materials provided with the distribution.
|
|
|
|
# 3. Neither the name of the copyright holder nor the names of its
|
|
# contributors may be used to endorse or promote products derived from
|
|
# this software without specific prior written permission.
|
|
|
|
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
|
|
|
|
import os
|
|
import sys
|
|
import argparse
|
|
from typing import List, Type, Tuple, Optional
|
|
import cuda.bindings.driver as cuda
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
import cutlass
|
|
import cutlass.cute as cute
|
|
import cutlass.cute.testing as testing
|
|
import cutlass.utils as utils
|
|
import cutlass.pipeline as pipeline
|
|
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
|
|
from cutlass.cute.nvgpu import cpasync, tcgen05
|
|
import cutlass.torch as cutlass_torch
|
|
import cutlass.utils.blackwell_helpers as sm100_utils
|
|
from cutlass.cute.runtime import from_dlpack
|
|
|
|
if __name__ == "__main__":
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
sys.path.insert(0, os.path.join(current_dir, "../.."))
|
|
|
|
from blackwell.mamba2_ssd.mamba2_ssd_reference import (
|
|
ssd_reference_fp32_all,
|
|
ssd_reference_lowprecision_intermediates,
|
|
analyze_relative_diffs,
|
|
)
|
|
from blackwell.mamba2_ssd.mamba2_ssd_tile_scheduler import (
|
|
Mamba2SSDTileSchedulerParams,
|
|
Mamba2SSDTileScheduler,
|
|
)
|
|
|
|
|
|
class SSDKernel:
|
|
def __init__(
|
|
self,
|
|
io_dtype: Type[cutlass.Numeric],
|
|
cumsum_delta_dtype: Type[cutlass.Numeric],
|
|
acc_dtype: Type[cutlass.Numeric],
|
|
L: int,
|
|
D: int,
|
|
N: int,
|
|
has_d: bool,
|
|
d_has_hdim: bool,
|
|
):
|
|
self.io_dtype: Type[cutlass.Numeric] = io_dtype
|
|
self.acc_dtype: Type[cutlass.Numeric] = acc_dtype
|
|
self.cumsum_delta_dtype: Type[cutlass.Numeric] = cumsum_delta_dtype
|
|
# has_d means epilog warp performs Y += X*D fusion
|
|
self.has_d: bool = has_d
|
|
# d_has_hdim = True means D is (D, EH) shape and loaded by TMA
|
|
# d_has_hdim = False means D is (1, EH) shape and loaded directly to register
|
|
self.d_has_hdim: bool = d_has_hdim
|
|
self.tile_shape = (L, D, N)
|
|
|
|
assert io_dtype in {
|
|
cutlass.Float16,
|
|
cutlass.BFloat16,
|
|
}, "Do not support other I/O types."
|
|
assert acc_dtype in {cutlass.Float32}, "Do not support other ACC types."
|
|
assert cumsum_delta_dtype in {cutlass.Float32}, (
|
|
"Do not support other cumsum types."
|
|
)
|
|
assert not (not has_d and d_has_hdim), "D cannot have Hdim if has_d is False"
|
|
|
|
# Hardcode default setting
|
|
self.use_2cta_instrs = False
|
|
self.cluster_shape_mnk = (1, 1, 1)
|
|
self.epi_tile = (128, 32)
|
|
|
|
# Setup mma tile shapes
|
|
self.tile_shape_mnk_intra1 = (L, L, N)
|
|
self.tile_shape_mnk_intra2 = (L, D, L)
|
|
self.tile_shape_mnk_inter1 = (N, D, L)
|
|
self.tile_shape_mnk_inter2 = (L, D, N)
|
|
|
|
self.cta_group = (
|
|
tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE
|
|
)
|
|
|
|
# Launch config
|
|
self.occupancy = 1
|
|
self.mma_inter_warp_id = 0
|
|
self.mma_intra_warp_id = 1
|
|
self.tma_b_c_warp_id = 2
|
|
self.tma_deltas_x_d_warp_id = 3
|
|
self.pre_inter_warp_id = [4, 5, 6, 7]
|
|
self.pre_intra_warp_id = [8, 9, 10, 11]
|
|
self.epilog_warp_id = [12, 13, 14, 15]
|
|
self.threads_per_cta = 32 * len(
|
|
(
|
|
self.mma_inter_warp_id,
|
|
self.mma_intra_warp_id,
|
|
self.tma_b_c_warp_id,
|
|
self.tma_deltas_x_d_warp_id,
|
|
*self.pre_inter_warp_id,
|
|
*self.pre_intra_warp_id,
|
|
*self.epilog_warp_id,
|
|
)
|
|
)
|
|
self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100")
|
|
|
|
# Named barriers
|
|
self.pre_inter_sync_barrier = pipeline.NamedBarrier(
|
|
barrier_id=1,
|
|
num_threads=len(self.pre_inter_warp_id) * 32,
|
|
)
|
|
self.epilog_sync_barrier = pipeline.NamedBarrier(
|
|
barrier_id=2,
|
|
num_threads=len(self.epilog_warp_id) * 32,
|
|
)
|
|
self.tmem_dealloc_sync_barrier = pipeline.NamedBarrier(
|
|
barrier_id=3,
|
|
num_threads=self.threads_per_cta,
|
|
)
|
|
|
|
# Number of registers used by each warp
|
|
self.num_regs_uniform_warps = 24
|
|
self.num_regs_pre_inter_warps = 168
|
|
self.num_regs_pre_intra_warps = 208
|
|
self.num_regs_epilogue_warps = 112
|
|
|
|
# Shared storage
|
|
self.shared_storage = None
|
|
|
|
# TMEM buffer offsets
|
|
self.tmem_intra1_acc_offset = 0
|
|
self.tmem_intra2_q_offset = 0
|
|
self.tmem_intra2_acc_offset = 0
|
|
self.tmem_inter1_acc_offset = 0
|
|
self.tmem_inter2_acc_offset = 0
|
|
self.num_tmem_cols_total = 0
|
|
|
|
def _setup_attributes(self):
|
|
(
|
|
tiled_mma_intra1,
|
|
tiled_mma_intra2,
|
|
tiled_mma_inter1,
|
|
tiled_mma_inter2,
|
|
) = self.make_tiled_mmas(
|
|
self.io_dtype,
|
|
self.acc_dtype,
|
|
self.cta_group,
|
|
self.tile_shape_mnk_intra1,
|
|
self.tile_shape_mnk_intra2,
|
|
self.tile_shape_mnk_inter1,
|
|
self.tile_shape_mnk_inter2,
|
|
)
|
|
|
|
self.cluster_layout_vmnk = cute.tiled_divide(
|
|
cute.make_layout(self.cluster_shape_mnk),
|
|
(tiled_mma_intra1.thr_id.shape,),
|
|
)
|
|
|
|
# Setup stages
|
|
(
|
|
self.input_stages,
|
|
self.output_stages,
|
|
self.internal_stages,
|
|
self.intra1_acc_stages,
|
|
) = self._compute_stages(
|
|
self.smem_capacity,
|
|
)
|
|
|
|
# Setup smem layouts
|
|
# X is B operand (from smem) of INTRA2_MMA and INTER1_MMA
|
|
self.x_smem_layout = sm100_utils.make_smem_layout_b(
|
|
tiled_mma_intra2,
|
|
self.tile_shape_mnk_intra2,
|
|
self.io_dtype,
|
|
self.input_stages,
|
|
)
|
|
self.num_x_load_bytes = cute.size_in_bytes(
|
|
self.io_dtype, cute.slice_(self.x_smem_layout, (None, None, None, 0))
|
|
)
|
|
|
|
# XT is same shape as ACC operand of INTER2_MMA, before postprocessing by EPILOG
|
|
self.xt_smem_layout = sm100_utils.make_smem_layout_epi(
|
|
self.io_dtype,
|
|
utils.LayoutEnum.COL_MAJOR,
|
|
self.tile_shape_mnk_intra2[:2],
|
|
self.input_stages,
|
|
)
|
|
|
|
# B is B operand (from smem) of INTRA1_MMA
|
|
self.b_smem_layout = sm100_utils.make_smem_layout_b(
|
|
tiled_mma_intra1,
|
|
self.tile_shape_mnk_intra1,
|
|
self.io_dtype,
|
|
self.input_stages,
|
|
)
|
|
self.num_b_load_bytes = cute.size_in_bytes(
|
|
self.io_dtype, cute.slice_(self.b_smem_layout, (None, None, None, 0))
|
|
)
|
|
|
|
# B_INTERNAL is also A operand (from smem) of INTER1_MMA, after preprocessed by PRE_INTER
|
|
self.bt_internal_smem_layout = sm100_utils.make_smem_layout_a(
|
|
tiled_mma_inter1,
|
|
self.tile_shape_mnk_inter1,
|
|
self.io_dtype,
|
|
self.internal_stages,
|
|
)
|
|
|
|
# B needs to be proprocessed to be used as A operand of INTER1_MMA
|
|
self.bt_smem_layout = cute.coalesce(
|
|
sm100_utils.make_smem_layout_epi(
|
|
self.io_dtype,
|
|
utils.LayoutEnum.ROW_MAJOR,
|
|
(self.tile_shape_mnk_inter1[0], self.tile_shape_mnk_inter1[2]),
|
|
self.input_stages,
|
|
),
|
|
target_profile=(1, 1, 1),
|
|
)
|
|
|
|
# C is A operand (from smem) of INTRA1_MMA and INTER2_MMA
|
|
self.c_smem_layout = sm100_utils.make_smem_layout_a(
|
|
tiled_mma_intra1,
|
|
self.tile_shape_mnk_intra1,
|
|
self.io_dtype,
|
|
self.input_stages,
|
|
)
|
|
self.num_c_load_bytes = cute.size_in_bytes(
|
|
self.io_dtype, cute.slice_(self.c_smem_layout, (None, None, None, 0))
|
|
)
|
|
|
|
# P is B operand (from smem) of INTER2_MMA, after preprocessed by PRE_INTER
|
|
self.p_smem_layout = sm100_utils.make_smem_layout_b(
|
|
tiled_mma_inter2,
|
|
self.tile_shape_mnk_inter2,
|
|
self.io_dtype,
|
|
self.internal_stages,
|
|
)
|
|
|
|
# PT is ACC operand (from tmem) of INTER1_MMA, after postprocessed by PRE_INTER
|
|
self.pt_smem_layout = sm100_utils.make_smem_layout_epi(
|
|
self.io_dtype,
|
|
utils.LayoutEnum.COL_MAJOR,
|
|
self.tile_shape_mnk_inter1[:2],
|
|
self.internal_stages,
|
|
)
|
|
|
|
# Q is A operand (from tmem) of INTRA2_MMA, after preprocessed by PRE_INTRA
|
|
self.q_tmem_layout = sm100_utils.make_smem_layout_a(
|
|
tiled_mma_intra2,
|
|
self.tile_shape_mnk_intra2,
|
|
self.io_dtype,
|
|
self.internal_stages,
|
|
)
|
|
|
|
# P is ACC operand (from tmem) of INTER1_MMA, to be TMA stored by PRE_INTER
|
|
self.p_smem_layout_store = sm100_utils.make_smem_layout_epi(
|
|
self.io_dtype,
|
|
utils.LayoutEnum.ROW_MAJOR,
|
|
self.tile_shape_mnk_inter2[1:],
|
|
self.internal_stages,
|
|
)
|
|
|
|
# Y is ACC operand (from smem) of INTER2_MMA and INTRA2_MMA, after postprocessed and TMA stored by EPILOG
|
|
self.y_smem_layout = sm100_utils.make_smem_layout_epi(
|
|
self.io_dtype,
|
|
utils.LayoutEnum.COL_MAJOR,
|
|
self.epi_tile,
|
|
self.output_stages,
|
|
)
|
|
|
|
# Delta is linear smem layouts for pre/post processing
|
|
self.delta_linear_smem_layout = cute.make_layout(
|
|
(self.tile_shape_mnk_inter1[2], self.input_stages)
|
|
)
|
|
self.num_delta_load_bytes = cute.size_in_bytes(
|
|
self.io_dtype, cute.slice_(self.delta_linear_smem_layout, (None, 0))
|
|
)
|
|
|
|
# Cumsum delta is linear smem layouts for pre/post processing
|
|
self.cumsum_delta_linear_smem_layout = cute.make_layout(
|
|
(self.tile_shape_mnk_inter1[2], self.input_stages)
|
|
)
|
|
self.num_cumsum_delta_load_bytes = cute.size_in_bytes(
|
|
self.cumsum_delta_dtype,
|
|
cute.slice_(self.cumsum_delta_linear_smem_layout, (None, 0)),
|
|
)
|
|
|
|
# D is linear smem layouts when d_has_hdim is True
|
|
self.d_linear_smem_layout = (
|
|
cute.make_layout((self.tile_shape_mnk_inter2[1], self.input_stages))
|
|
if self.d_has_hdim
|
|
else None
|
|
)
|
|
self.num_d_load_bytes = (
|
|
cute.size_in_bytes(
|
|
self.io_dtype,
|
|
cute.slice_(self.d_linear_smem_layout, (None, 0)),
|
|
)
|
|
if self.d_has_hdim
|
|
else 0
|
|
)
|
|
|
|
# Setup tmem offsets
|
|
(
|
|
self.tmem_intra1_acc_offset,
|
|
self.tmem_intra2_q_offset,
|
|
self.tmem_intra2_acc_offset,
|
|
self.tmem_inter1_acc_offset,
|
|
self.tmem_inter2_acc_offset,
|
|
self.num_tmem_cols_total,
|
|
) = self._plan_tmem_offsets(
|
|
tiled_mma_intra1,
|
|
self.tile_shape_mnk_intra1,
|
|
tiled_mma_intra2,
|
|
self.tile_shape_mnk_intra2,
|
|
tiled_mma_inter1,
|
|
self.tile_shape_mnk_inter1,
|
|
tiled_mma_inter2,
|
|
self.tile_shape_mnk_inter2,
|
|
self.internal_stages,
|
|
self.q_tmem_layout,
|
|
self.io_dtype,
|
|
self.internal_stages,
|
|
self.intra1_acc_stages,
|
|
)
|
|
|
|
return
|
|
|
|
@cute.jit
|
|
def __call__(
|
|
self,
|
|
x: cute.Tensor,
|
|
cumsum_delta: cute.Tensor,
|
|
delta: cute.Tensor,
|
|
b: cute.Tensor,
|
|
c: cute.Tensor,
|
|
y: cute.Tensor,
|
|
fstate: cute.Tensor,
|
|
d: cute.Tensor,
|
|
max_active_clusters: cutlass.Constexpr,
|
|
stream: cuda.CUstream,
|
|
):
|
|
self._setup_attributes()
|
|
(
|
|
tiled_mma_intra1,
|
|
tiled_mma_intra2,
|
|
tiled_mma_inter1,
|
|
tiled_mma_inter2,
|
|
) = self.make_tiled_mmas(
|
|
self.io_dtype,
|
|
self.acc_dtype,
|
|
self.cta_group,
|
|
self.tile_shape_mnk_intra1,
|
|
self.tile_shape_mnk_intra2,
|
|
self.tile_shape_mnk_inter1,
|
|
self.tile_shape_mnk_inter2,
|
|
)
|
|
|
|
# Setup TMA atoms and convert TMA tensors
|
|
# TMA load for A
|
|
x_op = sm100_utils.cluster_shape_to_tma_atom_B(
|
|
self.cluster_shape_mnk, tiled_mma_intra2.thr_id
|
|
)
|
|
tma_atom_x, tma_tensor_x = cute.nvgpu.make_tiled_tma_atom_B(
|
|
x_op,
|
|
x,
|
|
cute.slice_(self.x_smem_layout, (None, None, None, 0)),
|
|
self.tile_shape_mnk_intra2,
|
|
tiled_mma_intra2,
|
|
self.cluster_layout_vmnk.shape,
|
|
internal_type=(
|
|
cutlass.TFloat32 if x.element_type is cutlass.Float32 else None
|
|
),
|
|
)
|
|
|
|
# TMA load for B
|
|
b_op = sm100_utils.cluster_shape_to_tma_atom_B(
|
|
self.cluster_shape_mnk, tiled_mma_intra1.thr_id
|
|
)
|
|
tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B(
|
|
b_op,
|
|
b,
|
|
cute.slice_(self.b_smem_layout, (None, None, None, 0)),
|
|
self.tile_shape_mnk_intra1,
|
|
tiled_mma_intra1,
|
|
self.cluster_layout_vmnk.shape,
|
|
internal_type=(
|
|
cutlass.TFloat32 if b.element_type is cutlass.Float32 else None
|
|
),
|
|
)
|
|
|
|
# TMA load for C
|
|
c_op = sm100_utils.cluster_shape_to_tma_atom_A(
|
|
self.cluster_shape_mnk, tiled_mma_intra1.thr_id
|
|
)
|
|
tma_atom_c, tma_tensor_c = cute.nvgpu.make_tiled_tma_atom_A(
|
|
c_op,
|
|
c,
|
|
cute.slice_(self.c_smem_layout, (None, None, None, 0)),
|
|
self.tile_shape_mnk_intra1,
|
|
tiled_mma_intra1,
|
|
self.cluster_layout_vmnk.shape,
|
|
internal_type=(
|
|
cutlass.TFloat32 if c.element_type is cutlass.Float32 else None
|
|
),
|
|
)
|
|
|
|
# TMA load for delta
|
|
# TODO: use bulkcp instead of tma
|
|
delta_cta_v_layout = cute.slice_(
|
|
cute.make_identity_layout(delta.shape), (None, 0, 0, 0)
|
|
)
|
|
delta_linear_smem_layout = cute.slice_(self.delta_linear_smem_layout, (None, 0))
|
|
tma_atom_delta, tma_tensor_delta = cpasync.make_tiled_tma_atom(
|
|
cpasync.CopyBulkTensorTileG2SOp(),
|
|
delta,
|
|
delta_linear_smem_layout,
|
|
delta_cta_v_layout,
|
|
)
|
|
|
|
# TMA load for cumsum_delta
|
|
cumsum_delta_cta_v_layout = cute.slice_(
|
|
cute.make_identity_layout(cumsum_delta.shape), (None, 0, 0, 0)
|
|
)
|
|
cumsum_delta_linear_smem_layout = cute.slice_(
|
|
self.cumsum_delta_linear_smem_layout, (None, 0)
|
|
)
|
|
(
|
|
tma_atom_cumsum_delta,
|
|
tma_tensor_cumsum_delta,
|
|
) = cpasync.make_tiled_tma_atom(
|
|
cpasync.CopyBulkTensorTileG2SOp(),
|
|
cumsum_delta,
|
|
cumsum_delta_linear_smem_layout,
|
|
cumsum_delta_cta_v_layout,
|
|
)
|
|
|
|
tma_atom_d = None
|
|
tma_tensor_d = d
|
|
# TMA load for D
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
d_cta_v_layout = cute.slice_(cute.make_identity_layout(d.shape), (None, 0))
|
|
d_linear_smem_layout = cute.slice_(self.d_linear_smem_layout, (None, 0))
|
|
(
|
|
tma_atom_d,
|
|
tma_tensor_d,
|
|
) = cpasync.make_tiled_tma_atom(
|
|
cpasync.CopyBulkTensorTileG2SOp(),
|
|
d,
|
|
d_linear_smem_layout,
|
|
d_cta_v_layout,
|
|
)
|
|
|
|
# TMA store for y
|
|
y_smem_layout = cute.slice_(self.y_smem_layout, (None, None, 0))
|
|
tma_atom_y, tma_tensor_y = cpasync.make_tiled_tma_atom(
|
|
cpasync.CopyBulkTensorTileS2GOp(),
|
|
y,
|
|
y_smem_layout,
|
|
self.epi_tile,
|
|
)
|
|
|
|
# TMA store for fstate(p)
|
|
p_cta_v_layout = cute.slice_(
|
|
cute.make_identity_layout(fstate.shape), (None, None, 0, 0)
|
|
)
|
|
p_smem_layout_store = cute.slice_(self.p_smem_layout_store, (None, None, 0))
|
|
tma_atom_p, tma_tensor_p = cpasync.make_tiled_tma_atom(
|
|
cpasync.CopyBulkTensorTileS2GOp(),
|
|
fstate,
|
|
p_smem_layout_store,
|
|
p_cta_v_layout,
|
|
)
|
|
|
|
# Compute grid size
|
|
tile_sched_params, grid = self._compute_grid(y, b, max_active_clusters)
|
|
|
|
# Plan shared memory storage
|
|
swizzle_buffer_align_bytes = 1024
|
|
nonswizzle_buffer_align_bytes = 128
|
|
|
|
@cute.struct
|
|
class SharedStorage:
|
|
# Input stage barriers
|
|
x_full: cute.struct.MemRange[cutlass.Int64, self.input_stages] # type: ignore
|
|
x_empty: cute.struct.MemRange[cutlass.Int64, self.input_stages] # type: ignore
|
|
b_full: cute.struct.MemRange[cutlass.Int64, self.input_stages] # type: ignore
|
|
b_empty: cute.struct.MemRange[cutlass.Int64, self.input_stages] # type: ignore
|
|
c_full: cute.struct.MemRange[cutlass.Int64, self.input_stages] # type: ignore
|
|
c_empty: cute.struct.MemRange[cutlass.Int64, self.input_stages] # type: ignore
|
|
deltas_full: cute.struct.MemRange[cutlass.Int64, self.input_stages] # type: ignore
|
|
deltas_empty: cute.struct.MemRange[cutlass.Int64, self.input_stages] # type: ignore
|
|
d_full: cute.struct.MemRange[cutlass.Int64, self.input_stages] # type: ignore
|
|
d_empty: cute.struct.MemRange[cutlass.Int64, self.input_stages] # type: ignore
|
|
# Intra1 acc stage barriers
|
|
intra1_acc_full: cute.struct.MemRange[cutlass.Int64, self.intra1_acc_stages] # type: ignore
|
|
intra1_acc_empty: cute.struct.MemRange[
|
|
cutlass.Int64, self.intra1_acc_stages
|
|
] # type: ignore
|
|
# Internal stage barriers
|
|
intra2_q_full: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
|
|
intra2_q_empty: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
|
|
intra2_acc_full: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
|
|
intra2_acc_empty: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
|
|
inter1_b_full: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
|
|
inter1_b_empty: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
|
|
inter1_acc_full: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
|
|
inter1_acc_empty: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
|
|
inter2_p_full: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
|
|
inter2_p_empty: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
|
|
inter2_acc_full: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
|
|
inter2_acc_empty: cute.struct.MemRange[cutlass.Int64, self.internal_stages] # type: ignore
|
|
# Tmem holding buffer
|
|
tmem_holding_buf: cutlass.Int32
|
|
# Smem tensors
|
|
smem_x: cute.struct.Align[
|
|
cute.struct.MemRange[self.io_dtype, cute.cosize(self.x_smem_layout)],
|
|
swizzle_buffer_align_bytes,
|
|
]
|
|
smem_b: cute.struct.Align[
|
|
cute.struct.MemRange[self.io_dtype, cute.cosize(self.b_smem_layout)],
|
|
swizzle_buffer_align_bytes,
|
|
]
|
|
smem_bt_internal: cute.struct.Align[
|
|
cute.struct.MemRange[
|
|
self.io_dtype, cute.cosize(self.bt_internal_smem_layout)
|
|
],
|
|
swizzle_buffer_align_bytes,
|
|
]
|
|
smem_c: cute.struct.Align[
|
|
cute.struct.MemRange[self.io_dtype, cute.cosize(self.c_smem_layout)],
|
|
swizzle_buffer_align_bytes,
|
|
]
|
|
smem_p: cute.struct.Align[
|
|
cute.struct.MemRange[self.io_dtype, cute.cosize(self.p_smem_layout)],
|
|
swizzle_buffer_align_bytes,
|
|
]
|
|
smem_y: cute.struct.Align[
|
|
cute.struct.MemRange[self.io_dtype, cute.cosize(self.y_smem_layout)],
|
|
swizzle_buffer_align_bytes,
|
|
]
|
|
smem_cumsum_delta: cute.struct.Align[
|
|
cute.struct.MemRange[
|
|
self.cumsum_delta_dtype,
|
|
cute.cosize(self.cumsum_delta_linear_smem_layout),
|
|
],
|
|
nonswizzle_buffer_align_bytes,
|
|
]
|
|
smem_delta: cute.struct.Align[
|
|
cute.struct.MemRange[
|
|
self.io_dtype, cute.cosize(self.delta_linear_smem_layout)
|
|
],
|
|
nonswizzle_buffer_align_bytes,
|
|
]
|
|
smem_d: cute.struct.Align[
|
|
cute.struct.MemRange[
|
|
self.io_dtype,
|
|
cute.cosize(self.d_linear_smem_layout) if self.d_has_hdim else 0,
|
|
],
|
|
nonswizzle_buffer_align_bytes,
|
|
]
|
|
|
|
self.shared_storage = SharedStorage
|
|
if cutlass.const_expr(self.shared_storage.size_in_bytes() > self.smem_capacity):
|
|
raise ValueError(
|
|
f"SharedStorage size {self.shared_storage.size_in_bytes()} exceeds smem_capacity {self.smem_capacity}"
|
|
)
|
|
|
|
# Launch the kernel synchronously
|
|
self.kernel(
|
|
tma_atom_x,
|
|
tma_tensor_x,
|
|
tma_atom_b,
|
|
tma_tensor_b,
|
|
tma_atom_c,
|
|
tma_tensor_c,
|
|
tma_atom_p,
|
|
tma_tensor_p,
|
|
tma_atom_y,
|
|
tma_tensor_y,
|
|
tma_atom_delta,
|
|
tma_tensor_delta,
|
|
tma_atom_cumsum_delta,
|
|
tma_tensor_cumsum_delta,
|
|
tma_atom_d,
|
|
tma_tensor_d,
|
|
self.cluster_layout_vmnk,
|
|
self.x_smem_layout,
|
|
self.xt_smem_layout,
|
|
self.b_smem_layout,
|
|
self.bt_smem_layout,
|
|
self.bt_internal_smem_layout,
|
|
self.c_smem_layout,
|
|
self.pt_smem_layout,
|
|
self.p_smem_layout,
|
|
self.q_tmem_layout,
|
|
self.p_smem_layout_store,
|
|
self.y_smem_layout,
|
|
self.delta_linear_smem_layout,
|
|
self.cumsum_delta_linear_smem_layout,
|
|
self.d_linear_smem_layout,
|
|
self.epi_tile,
|
|
tile_sched_params,
|
|
).launch(
|
|
grid=grid,
|
|
block=[self.threads_per_cta, 1, 1],
|
|
cluster=self.cluster_shape_mnk,
|
|
min_blocks_per_mp=1,
|
|
stream=stream,
|
|
)
|
|
|
|
# GPU device kernel
|
|
@cute.kernel
|
|
def kernel(
|
|
self,
|
|
tma_atom_x: cute.CopyAtom,
|
|
tma_tensor_x: cute.Tensor,
|
|
tma_atom_b: cute.CopyAtom,
|
|
tma_tensor_b: cute.Tensor,
|
|
tma_atom_c: cute.CopyAtom,
|
|
tma_tensor_c: cute.Tensor,
|
|
tma_atom_p: cute.CopyAtom,
|
|
tma_tensor_p: cute.Tensor,
|
|
tma_atom_y: cute.CopyAtom,
|
|
tma_tensor_y: cute.Tensor,
|
|
tma_atom_delta: cute.CopyAtom,
|
|
tma_tensor_delta: cute.Tensor,
|
|
tma_atom_cumsum_delta: cute.CopyAtom,
|
|
tma_tensor_cumsum_delta: cute.Tensor,
|
|
tma_atom_d: Optional[cute.CopyAtom],
|
|
tma_tensor_d: cute.Tensor,
|
|
cluster_layout_vmnk: cute.Layout,
|
|
x_smem_layout: cute.ComposedLayout,
|
|
xt_smem_layout: cute.ComposedLayout,
|
|
b_smem_layout: cute.ComposedLayout,
|
|
bt_smem_layout: cute.ComposedLayout,
|
|
bt_internal_smem_layout: cute.ComposedLayout,
|
|
c_smem_layout: cute.ComposedLayout,
|
|
pt_smem_layout: cute.ComposedLayout,
|
|
p_smem_layout: cute.ComposedLayout,
|
|
q_tmem_layout: cute.ComposedLayout,
|
|
p_smem_layout_store: cute.ComposedLayout,
|
|
y_smem_layout: cute.ComposedLayout,
|
|
delta_linear_smem_layout: cute.Layout,
|
|
cumsum_delta_linear_smem_layout: cute.Layout,
|
|
d_linear_smem_layout: Optional[cute.Layout],
|
|
epi_tile: cute.Tile,
|
|
tile_sched_params: Mamba2SSDTileSchedulerParams,
|
|
):
|
|
warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
|
|
|
|
# Prefetch tma descriptor
|
|
if warp_idx == 0:
|
|
tma_atoms = [
|
|
tma_atom_x,
|
|
tma_atom_b,
|
|
tma_atom_c,
|
|
tma_atom_p,
|
|
tma_atom_y,
|
|
tma_atom_delta,
|
|
tma_atom_cumsum_delta,
|
|
]
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
tma_atoms.append(tma_atom_d)
|
|
for tma_atom in tma_atoms:
|
|
cpasync.prefetch_descriptor(tma_atom)
|
|
|
|
# Static consts
|
|
D = cute.size(tma_tensor_x, mode=[0])
|
|
L = cute.size(tma_tensor_x, mode=[1])
|
|
N = cute.size(tma_tensor_b, mode=[1])
|
|
# Dynamic values
|
|
C = cute.size(tma_tensor_x, mode=[2])
|
|
EH = cute.size(tma_tensor_x, mode=[3])
|
|
B = cute.size(tma_tensor_x, mode=[4])
|
|
G = cute.size(tma_tensor_b, mode=[3])
|
|
NGROUP_RATIO = EH // G
|
|
|
|
# Make TiledMma
|
|
(
|
|
tiled_mma_intra1,
|
|
tiled_mma_intra2,
|
|
tiled_mma_inter1,
|
|
tiled_mma_inter2,
|
|
) = self.make_tiled_mmas(
|
|
self.io_dtype,
|
|
self.acc_dtype,
|
|
self.cta_group,
|
|
self.tile_shape_mnk_intra1,
|
|
self.tile_shape_mnk_intra2,
|
|
self.tile_shape_mnk_inter1,
|
|
self.tile_shape_mnk_inter2,
|
|
)
|
|
|
|
# Setup cta/thread coordinates
|
|
# Block coord
|
|
bidx, bidy, bidz = cute.arch.block_idx()
|
|
mma_tile_coord_v = bidx % cute.size(tiled_mma_intra1.thr_id.shape)
|
|
cta_rank_in_cluster = cute.arch.make_warp_uniform(
|
|
cute.arch.block_idx_in_cluster()
|
|
)
|
|
block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(
|
|
cta_rank_in_cluster
|
|
)
|
|
# Workload coord
|
|
tile_sched = Mamba2SSDTileScheduler.create(
|
|
tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim()
|
|
)
|
|
work_tile = tile_sched.initial_work_tile_info()
|
|
|
|
# Thread/warp coord
|
|
tidx, _, _ = cute.arch.thread_idx()
|
|
# Thread coord inside specialized warps
|
|
local_tidx = tidx % 128
|
|
local_warp_idx = cute.arch.make_warp_uniform(local_tidx // 32)
|
|
|
|
# Alloc and init smem tensors and pipelines
|
|
smem = utils.SmemAllocator()
|
|
smem_storage = smem.allocate(self.shared_storage)
|
|
|
|
# Setup smem tensors
|
|
smem_x = smem_storage.smem_x.get_tensor(
|
|
x_smem_layout.outer, swizzle=x_smem_layout.inner
|
|
)
|
|
smem_xt = smem_storage.smem_x.get_tensor(
|
|
xt_smem_layout.outer, swizzle=xt_smem_layout.inner
|
|
)
|
|
smem_b = smem_storage.smem_b.get_tensor(
|
|
b_smem_layout.outer, swizzle=b_smem_layout.inner
|
|
)
|
|
smem_bt = smem_storage.smem_b.get_tensor(
|
|
bt_smem_layout.outer, swizzle=bt_smem_layout.inner
|
|
)
|
|
smem_bt_internal = smem_storage.smem_bt_internal.get_tensor(
|
|
bt_internal_smem_layout.outer, swizzle=bt_internal_smem_layout.inner
|
|
)
|
|
smem_c = smem_storage.smem_c.get_tensor(
|
|
c_smem_layout.outer, swizzle=c_smem_layout.inner
|
|
)
|
|
smem_p = smem_storage.smem_p.get_tensor(
|
|
p_smem_layout.outer, swizzle=p_smem_layout.inner
|
|
)
|
|
smem_pt = smem_storage.smem_p.get_tensor(
|
|
pt_smem_layout.outer, swizzle=pt_smem_layout.inner
|
|
)
|
|
smem_p_store = smem_storage.smem_p.get_tensor(
|
|
p_smem_layout_store.outer, swizzle=p_smem_layout_store.inner
|
|
)
|
|
smem_y = smem_storage.smem_y.get_tensor(
|
|
y_smem_layout.outer, swizzle=y_smem_layout.inner
|
|
)
|
|
smem_cumsum_delta = smem_storage.smem_cumsum_delta.get_tensor(
|
|
cumsum_delta_linear_smem_layout
|
|
)
|
|
smem_delta = smem_storage.smem_delta.get_tensor(delta_linear_smem_layout)
|
|
smem_d = None
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
smem_d = smem_storage.smem_d.get_tensor(d_linear_smem_layout)
|
|
|
|
# Init mbarrier for pipeline
|
|
x_pipeline = self.make_and_init_x_pipeline(smem_storage.x_full.data_ptr())
|
|
b_pipeline = self.make_and_init_b_pipeline(smem_storage.b_full.data_ptr())
|
|
c_pipeline = self.make_and_init_c_pipeline(smem_storage.c_full.data_ptr())
|
|
deltas_pipeline = self.make_and_init_deltas_pipeline(
|
|
smem_storage.deltas_full.data_ptr()
|
|
)
|
|
d_pipeline = self.make_and_init_d_pipeline(smem_storage.d_full.data_ptr())
|
|
intra1_acc_pipeline = self.make_and_init_intra1_acc_pipeline(
|
|
smem_storage.intra1_acc_full.data_ptr()
|
|
)
|
|
intra2_q_pipeline = self.make_and_init_intra2_q_pipeline(
|
|
smem_storage.intra2_q_full.data_ptr()
|
|
)
|
|
intra2_acc_pipeline = self.make_and_init_intra2_acc_pipeline(
|
|
smem_storage.intra2_acc_full.data_ptr()
|
|
)
|
|
inter1_b_pipeline = self.make_and_init_inter1_b_pipeline(
|
|
smem_storage.inter1_b_full.data_ptr()
|
|
)
|
|
inter1_acc_pipeline = self.make_and_init_inter1_acc_pipeline(
|
|
smem_storage.inter1_acc_full.data_ptr()
|
|
)
|
|
inter2_p_pipeline = self.make_and_init_inter2_p_pipeline(
|
|
smem_storage.inter2_p_full.data_ptr()
|
|
)
|
|
inter2_acc_pipeline = self.make_and_init_inter2_acc_pipeline(
|
|
smem_storage.inter2_acc_full.data_ptr()
|
|
)
|
|
|
|
# Cluster arrive after barrier init
|
|
pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mnk, is_relaxed=True)
|
|
|
|
# Cluster wait before tmem alloc
|
|
pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mnk)
|
|
|
|
tmem_alloc_barrier = pipeline.NamedBarrier(
|
|
barrier_id=0,
|
|
num_threads=self.threads_per_cta,
|
|
)
|
|
tmem = utils.TmemAllocator(
|
|
smem_storage.tmem_holding_buf,
|
|
barrier_for_retrieve=tmem_alloc_barrier,
|
|
allocator_warp_id=self.epilog_warp_id[0],
|
|
)
|
|
tmem.allocate(self.num_tmem_cols_total)
|
|
|
|
# Barrier before retrieve tensor memory ptr from shared memory
|
|
tmem.wait_for_alloc()
|
|
|
|
# Retrieve tmem ptr
|
|
tmem_ptr_base = tmem.retrieve_ptr(self.acc_dtype)
|
|
|
|
# Specialized TMA load Delta/CumsumDelta/X warp
|
|
if warp_idx == self.tma_deltas_x_d_warp_id:
|
|
# Dealloc regs for pre-inter/pre-intra warps
|
|
cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps)
|
|
|
|
# ((ATOM_V, REST_V), INPUT_STAGE)
|
|
# ((ATOM_V, REST_V), 1, 1, C, EH, B)
|
|
tXsX, tXgX_pre_slice = self.tma_partition_for_mma_b_operand(
|
|
tma_atom_x,
|
|
tma_tensor_x,
|
|
smem_x,
|
|
tiled_mma_intra2,
|
|
cluster_layout_vmnk,
|
|
mma_tile_coord_v,
|
|
block_in_cluster_coord_vmnk,
|
|
)
|
|
|
|
# ((ATOM_V, REST_V), INPUT_STAGE)
|
|
# ((ATOM_V, REST_V), 1, C, EH, B)
|
|
tDeltasDelta, tDeltagDelta_pre_slice = self.tma_partition_with_shape(
|
|
tma_atom_delta,
|
|
tma_tensor_delta,
|
|
smem_delta,
|
|
(self.tile_shape_mnk_inter1[2],),
|
|
)
|
|
# ((ATOM_V, REST_V), INPUT_STAGE)
|
|
# ((ATOM_V, REST_V), 1, C, EH, B)
|
|
(
|
|
tDeltasCumsumDelta,
|
|
tDeltagCumsumDelta_pre_slice,
|
|
) = self.tma_partition_with_shape(
|
|
tma_atom_cumsum_delta,
|
|
tma_tensor_cumsum_delta,
|
|
smem_cumsum_delta,
|
|
(self.tile_shape_mnk_inter1[2],),
|
|
)
|
|
|
|
tDsD = None
|
|
tDgD_pre_slice = None
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
# Partition global/shared tensor for D
|
|
# ((ATOM_V, REST_V), INPUT_STAGE)
|
|
# ((ATOM_V, REST_V), 1, EH)
|
|
tDsD, tDgD_pre_slice = self.tma_partition_with_shape(
|
|
tma_atom_d, tma_tensor_d, smem_d, (self.tile_shape_mnk_inter2[1],)
|
|
)
|
|
|
|
# Pipeline X/Delta/CumsumDelta/D producer state
|
|
x_producer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Producer, self.input_stages
|
|
)
|
|
deltas_producer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Producer, self.input_stages
|
|
)
|
|
d_producer_state = None
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
# D is loaded by TMA only when d_has_hdim is True
|
|
d_producer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Producer, self.input_stages
|
|
)
|
|
|
|
while work_tile.is_valid_tile:
|
|
b_idx, eh_idx, g_idx = work_tile.tile_idx
|
|
|
|
# Slice global tensor to current tile idx
|
|
# ((ATOM_V, REST_V), C)
|
|
tXgX = tXgX_pre_slice[None, 0, 0, None, eh_idx, b_idx]
|
|
tDeltagDelta = tDeltagDelta_pre_slice[None, 0, None, eh_idx, b_idx]
|
|
tDeltagCumsumDelta = tDeltagCumsumDelta_pre_slice[
|
|
None, 0, None, eh_idx, b_idx
|
|
]
|
|
tDgD = None
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
# ((ATOM_V, REST_V))
|
|
tDgD = tDgD_pre_slice[None, 0, eh_idx]
|
|
|
|
# Reset count for pipeline state
|
|
x_producer_state.reset_count()
|
|
deltas_producer_state.reset_count()
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
d_producer_state.reset_count()
|
|
|
|
# Peek (try_wait) X/deltas buffer empty status
|
|
peek_x_empty_status = self.conditional_producer_try_acquire(
|
|
x_producer_state, x_pipeline, C
|
|
)
|
|
peek_deltas_empty_status = self.conditional_producer_try_acquire(
|
|
deltas_producer_state, deltas_pipeline, C
|
|
)
|
|
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
# Wait for D buffer empty
|
|
d_pipeline.producer_acquire(d_producer_state)
|
|
# TMA load D
|
|
cute.copy(
|
|
tma_atom_d,
|
|
tDgD,
|
|
tDsD[None, d_producer_state.index],
|
|
tma_bar_ptr=d_pipeline.producer_get_barrier(d_producer_state),
|
|
)
|
|
# Advance D producer state
|
|
d_producer_state.advance()
|
|
|
|
# Batched load over C dimension
|
|
for chunk_idx in cutlass.range(C, unroll=1):
|
|
# Conditionally wait for X buffer empty
|
|
x_pipeline.producer_acquire(x_producer_state, peek_x_empty_status)
|
|
|
|
# TMA load X
|
|
cute.copy(
|
|
tma_atom_x,
|
|
tXgX[None, x_producer_state.count],
|
|
tXsX[None, x_producer_state.index],
|
|
tma_bar_ptr=x_pipeline.producer_get_barrier(x_producer_state),
|
|
)
|
|
|
|
# Conditionally wait for deltas buffer empty
|
|
deltas_pipeline.producer_acquire(
|
|
deltas_producer_state, peek_deltas_empty_status
|
|
)
|
|
|
|
# TMA load Delta/CumsumDelta
|
|
cute.copy(
|
|
tma_atom_delta,
|
|
tDeltagDelta[None, deltas_producer_state.count],
|
|
tDeltasDelta[None, deltas_producer_state.index],
|
|
tma_bar_ptr=deltas_pipeline.producer_get_barrier(
|
|
deltas_producer_state
|
|
),
|
|
)
|
|
cute.copy(
|
|
tma_atom_cumsum_delta,
|
|
tDeltagCumsumDelta[None, deltas_producer_state.count],
|
|
tDeltasCumsumDelta[None, deltas_producer_state.index],
|
|
tma_bar_ptr=deltas_pipeline.producer_get_barrier(
|
|
deltas_producer_state
|
|
),
|
|
)
|
|
|
|
# Advance X/deltas producer state
|
|
x_producer_state.advance()
|
|
deltas_producer_state.advance()
|
|
|
|
# Peek (try_wait) X/deltas buffer empty status
|
|
peek_x_empty_status = self.conditional_producer_try_acquire(
|
|
x_producer_state, x_pipeline, C
|
|
)
|
|
peek_deltas_empty_status = self.conditional_producer_try_acquire(
|
|
deltas_producer_state, deltas_pipeline, C
|
|
)
|
|
# END of for chunk_idx in cutlass.range(C, unroll=1)
|
|
|
|
# Advance to next tile
|
|
tile_sched.advance_to_next_work()
|
|
work_tile = tile_sched.get_current_work()
|
|
# END of while work_tile.is_valid_tile
|
|
|
|
# Producer tail for X/Deltas/D
|
|
x_pipeline.producer_tail(x_producer_state)
|
|
deltas_pipeline.producer_tail(deltas_producer_state)
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
d_pipeline.producer_tail(d_producer_state)
|
|
# END of specialized tma load X/Deltas/D warp
|
|
|
|
# Specialized TMA load B/C warp
|
|
elif warp_idx == self.tma_b_c_warp_id:
|
|
# Dealloc regs for pre-inter/pre-intra warps
|
|
cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps)
|
|
|
|
# ((ATOM_V, REST_V), INPUT_STAGE)
|
|
# ((ATOM_V, REST_V), 1, 1, C, G, B)
|
|
tBsB, tBgB_pre_slice = self.tma_partition_for_mma_b_operand(
|
|
tma_atom_b,
|
|
tma_tensor_b,
|
|
smem_b,
|
|
tiled_mma_intra1,
|
|
cluster_layout_vmnk,
|
|
mma_tile_coord_v,
|
|
block_in_cluster_coord_vmnk,
|
|
)
|
|
|
|
# ((ATOM_V, REST_V), INPUT_STAGE)
|
|
# ((ATOM_V, REST_V), 1, 1, C, G, B)
|
|
tCsC, tCgC_pre_slice = self.tma_partition_for_mma_a_operand(
|
|
tma_atom_c,
|
|
tma_tensor_c,
|
|
smem_c,
|
|
tiled_mma_intra1,
|
|
cluster_layout_vmnk,
|
|
mma_tile_coord_v,
|
|
block_in_cluster_coord_vmnk,
|
|
)
|
|
|
|
# Pipeline B/C producer state
|
|
b_producer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Producer, self.input_stages
|
|
)
|
|
c_producer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Producer, self.input_stages
|
|
)
|
|
|
|
while work_tile.is_valid_tile:
|
|
b_idx, eh_idx, g_idx = work_tile.tile_idx
|
|
|
|
# Slice global tensor to current tile idx
|
|
# ((ATOM_V, REST_V), C)
|
|
tBgB = tBgB_pre_slice[None, 0, 0, None, g_idx, b_idx]
|
|
tCgC = tCgC_pre_slice[None, 0, 0, None, g_idx, b_idx]
|
|
|
|
# Reset count for pipeline state
|
|
b_producer_state.reset_count()
|
|
c_producer_state.reset_count()
|
|
|
|
# Peek (try_wait) B/C buffer empty status
|
|
peek_b_empty_status = self.conditional_producer_try_acquire(
|
|
b_producer_state, b_pipeline, C
|
|
)
|
|
peek_c_empty_status = self.conditional_producer_try_acquire(
|
|
c_producer_state, c_pipeline, C
|
|
)
|
|
|
|
# Batched load over C dimension
|
|
for chunk_idx in cutlass.range(C, unroll=1):
|
|
# Conditionally wait for B buffer empty
|
|
b_pipeline.producer_acquire(b_producer_state, peek_b_empty_status)
|
|
|
|
# TMA load B
|
|
cute.copy(
|
|
tma_atom_b,
|
|
tBgB[None, b_producer_state.count],
|
|
tBsB[None, b_producer_state.index],
|
|
tma_bar_ptr=b_pipeline.producer_get_barrier(b_producer_state),
|
|
)
|
|
|
|
# Conditionally wait for C buffer empty
|
|
c_pipeline.producer_acquire(c_producer_state, peek_c_empty_status)
|
|
|
|
# TMA load C
|
|
cute.copy(
|
|
tma_atom_c,
|
|
tCgC[None, c_producer_state.count],
|
|
tCsC[None, c_producer_state.index],
|
|
tma_bar_ptr=c_pipeline.producer_get_barrier(c_producer_state),
|
|
)
|
|
|
|
# Advance B/C producer state
|
|
b_producer_state.advance()
|
|
c_producer_state.advance()
|
|
|
|
# Peek (try_wait) B/C buffer empty status
|
|
peek_b_empty_status = self.conditional_producer_try_acquire(
|
|
b_producer_state, b_pipeline, C
|
|
)
|
|
peek_c_empty_status = self.conditional_producer_try_acquire(
|
|
c_producer_state, c_pipeline, C
|
|
)
|
|
# END of for chunk_idx in cutlass.range(C, unroll=1)
|
|
|
|
# Advance to next tile
|
|
tile_sched.advance_to_next_work()
|
|
work_tile = tile_sched.get_current_work()
|
|
# END of while work_tile.is_valid_tile
|
|
|
|
# Producer tail for B/C
|
|
b_pipeline.producer_tail(b_producer_state)
|
|
c_pipeline.producer_tail(c_producer_state)
|
|
# END of specialized tma load B/C warp
|
|
|
|
# Specialized MMA Intra warp
|
|
elif warp_idx == self.mma_intra_warp_id:
|
|
# Dealloc regs for pre-inter/pre-intra warps
|
|
cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps)
|
|
|
|
# Make shared/tmem fragments for INTRA_MMA1 B/C/ACC
|
|
# (MMA, MMA_N, MMA_K, INPUT_STAGE)
|
|
# (MMA, MMA_M, MMA_K, INPUT_STAGE)
|
|
# (MMA, MMA_M, MMA_N, INTRA1_ACC_STAGE)
|
|
tCrC, tCrB, tCtAccIntra1 = self.mma_partition_ss(
|
|
tiled_mma_intra1,
|
|
self.tile_shape_mnk_intra1,
|
|
smem_c,
|
|
smem_b,
|
|
tmem_ptr_base + self.tmem_intra1_acc_offset,
|
|
self.intra1_acc_stages,
|
|
)
|
|
|
|
# Make shared/tmem fragments for INTRA_MMA2 X/Q/ACC
|
|
# (MMA, MMA_M, MMA_K, INTERNAL_STAGE)
|
|
# (MMA, MMA_N, MMA_K, INPUT_STAGE)
|
|
# (MMA, MMA_M, MMA_N, INTERNAL_STAGE)
|
|
tCrQ, tCrX, tCtAccIntra2 = self.mma_partition_ts(
|
|
tiled_mma_intra2,
|
|
self.tile_shape_mnk_intra2,
|
|
q_tmem_layout,
|
|
smem_x,
|
|
tmem_ptr_base + self.tmem_intra2_q_offset,
|
|
tmem_ptr_base + self.tmem_intra2_acc_offset,
|
|
self.internal_stages,
|
|
)
|
|
|
|
# Pipeline B/C/X/INTRA2_Q consumer state
|
|
b_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.input_stages
|
|
)
|
|
c_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.input_stages
|
|
)
|
|
x_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.input_stages
|
|
)
|
|
intra2_q_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.internal_stages
|
|
)
|
|
|
|
# Pipeline INTRA1_ACC/INTRA2_ACC producer state
|
|
intra1_acc_producer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Producer, self.intra1_acc_stages
|
|
)
|
|
intra2_acc_producer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Producer, self.internal_stages
|
|
)
|
|
|
|
while work_tile.is_valid_tile:
|
|
# Reset count for pipeline state
|
|
b_consumer_state.reset_count()
|
|
c_consumer_state.reset_count()
|
|
intra1_acc_producer_state.reset_count()
|
|
x_consumer_state.reset_count()
|
|
intra2_q_consumer_state.reset_count()
|
|
intra2_acc_producer_state.reset_count()
|
|
|
|
# Peek (try_wait) B/C/X/INTRA1_ACC buffer full/full/full/empty status
|
|
peek_b_full_status = self.conditional_consumer_try_wait(
|
|
b_consumer_state, b_pipeline, C
|
|
)
|
|
peek_c_full_status = self.conditional_consumer_try_wait(
|
|
c_consumer_state, c_pipeline, C
|
|
)
|
|
peek_wr_intra1_acc_empty_status = self.conditional_producer_try_acquire(
|
|
intra1_acc_producer_state, intra1_acc_pipeline, C
|
|
)
|
|
peek_x_full_status = self.conditional_consumer_try_wait(
|
|
x_consumer_state, x_pipeline, C
|
|
)
|
|
|
|
# Manual pipeline: unrolled INTRA_MMA1 chunk_idx = 0 loop
|
|
# Conditionally wait for B/C/INTRA1_ACC buffer full/full/empty
|
|
b_pipeline.consumer_wait(b_consumer_state, peek_b_full_status)
|
|
c_pipeline.consumer_wait(c_consumer_state, peek_c_full_status)
|
|
intra1_acc_pipeline.producer_acquire(
|
|
intra1_acc_producer_state, peek_wr_intra1_acc_empty_status
|
|
)
|
|
|
|
# INTRA_MMA1
|
|
tiled_mma_intra1 = self.exec_mma(
|
|
tiled_mma_intra1,
|
|
tCtAccIntra1,
|
|
tCrC,
|
|
tCrB,
|
|
intra1_acc_producer_state,
|
|
c_consumer_state,
|
|
b_consumer_state,
|
|
)
|
|
|
|
# Async arrive B/C/INTRA1_ACC buffer empty/empty/full
|
|
b_pipeline.consumer_release(
|
|
b_consumer_state, pipeline.PipelineOp.TCGen05Mma
|
|
)
|
|
c_pipeline.consumer_release(c_consumer_state)
|
|
intra1_acc_pipeline.producer_commit(intra1_acc_producer_state)
|
|
|
|
# Advance B/C/INTRA1_ACC state
|
|
b_consumer_state.advance()
|
|
c_consumer_state.advance()
|
|
intra1_acc_producer_state.advance()
|
|
|
|
# Peek (try_wait) B/C/INTRA1_ACC buffer full/full/empty for chunk_idx = chunk_idx + 1
|
|
peek_b_full_status = self.conditional_consumer_try_wait(
|
|
b_consumer_state, b_pipeline, C
|
|
)
|
|
peek_c_full_status = self.conditional_consumer_try_wait(
|
|
c_consumer_state, c_pipeline, C
|
|
)
|
|
peek_wr_intra1_acc_empty_status = self.conditional_producer_try_acquire(
|
|
intra1_acc_producer_state, intra1_acc_pipeline, C
|
|
)
|
|
|
|
# Manual pipeline: batched gemm over C-1 dimension
|
|
for chunk_idx in cutlass.range(C - 1, unroll=1):
|
|
# Conditionally wait for B/C/INTRA1_ACC buffer full/full/empty
|
|
b_pipeline.consumer_wait(b_consumer_state, peek_b_full_status)
|
|
c_pipeline.consumer_wait(c_consumer_state, peek_c_full_status)
|
|
intra1_acc_pipeline.producer_acquire(
|
|
intra1_acc_producer_state, peek_wr_intra1_acc_empty_status
|
|
)
|
|
|
|
# INTRA_MMA1
|
|
tiled_mma_intra1 = self.exec_mma(
|
|
tiled_mma_intra1,
|
|
tCtAccIntra1,
|
|
tCrC,
|
|
tCrB,
|
|
intra1_acc_producer_state,
|
|
c_consumer_state,
|
|
b_consumer_state,
|
|
)
|
|
|
|
# Async arrive B/C/INTRA1_ACC buffer empty/empty/full
|
|
b_pipeline.consumer_release(
|
|
b_consumer_state, pipeline.PipelineOp.TCGen05Mma
|
|
)
|
|
c_pipeline.consumer_release(c_consumer_state)
|
|
intra1_acc_pipeline.producer_commit(intra1_acc_producer_state)
|
|
|
|
# Conditionally wait for X/INTRA2_Q/INTRA2_ACC buffer full/full/empty
|
|
x_pipeline.consumer_wait(x_consumer_state, peek_x_full_status)
|
|
intra2_q_pipeline.consumer_wait(intra2_q_consumer_state)
|
|
intra2_acc_pipeline.producer_acquire(intra2_acc_producer_state)
|
|
|
|
# INTRA_MMA2
|
|
tiled_mma_intra2 = self.exec_mma(
|
|
tiled_mma_intra2,
|
|
tCtAccIntra2,
|
|
tCrQ,
|
|
tCrX,
|
|
intra2_acc_producer_state,
|
|
intra2_q_consumer_state,
|
|
x_consumer_state,
|
|
)
|
|
|
|
# Async arrive X/INTRA2_Q/INTRA2_ACC buffer empty/empty/full
|
|
if cutlass.const_expr(self.has_d):
|
|
x_pipeline.consumer_release(
|
|
x_consumer_state, pipeline.PipelineOp.TCGen05Mma
|
|
)
|
|
else:
|
|
x_pipeline.consumer_release(x_consumer_state)
|
|
intra2_q_pipeline.consumer_release(intra2_q_consumer_state)
|
|
intra2_acc_pipeline.producer_commit(intra2_acc_producer_state)
|
|
|
|
# Advance B/C/INTRA1_ACC cstate
|
|
b_consumer_state.advance()
|
|
c_consumer_state.advance()
|
|
intra1_acc_producer_state.advance()
|
|
|
|
# Peek (try_wait) B/C/INTRA1_ACC buffer full/full/empty for chunk_idx = chunk_idx + 1
|
|
peek_b_full_status = self.conditional_consumer_try_wait(
|
|
b_consumer_state, b_pipeline, C
|
|
)
|
|
peek_c_full_status = self.conditional_consumer_try_wait(
|
|
c_consumer_state, c_pipeline, C
|
|
)
|
|
peek_wr_intra1_acc_empty_status = (
|
|
self.conditional_producer_try_acquire(
|
|
intra1_acc_producer_state, intra1_acc_pipeline, C
|
|
)
|
|
)
|
|
|
|
# Advance X/INTRA2_Q/INTRA2_ACC state
|
|
x_consumer_state.advance()
|
|
intra2_q_consumer_state.advance()
|
|
intra2_acc_producer_state.advance()
|
|
|
|
# Peek (try_wait) X buffer full for chunk_idx = chunk_idx + 1
|
|
peek_x_full_status = self.conditional_consumer_try_wait(
|
|
x_consumer_state, x_pipeline, C
|
|
)
|
|
# END of for chunk_idx in cutlass.range(C-1, unroll=1)
|
|
|
|
# Manual pipeline: unrolled INTRA_MMA2 chunk_idx = C-1 loop
|
|
# Conditionally wait for X/INTRA2_Q/INTRA2_ACC buffer full/full/empty
|
|
x_pipeline.consumer_wait(x_consumer_state, peek_x_full_status)
|
|
intra2_q_pipeline.consumer_wait(intra2_q_consumer_state)
|
|
intra2_acc_pipeline.producer_acquire(intra2_acc_producer_state)
|
|
|
|
# INTRA_MMA2
|
|
tiled_mma_intra2 = self.exec_mma(
|
|
tiled_mma_intra2,
|
|
tCtAccIntra2,
|
|
tCrQ,
|
|
tCrX,
|
|
intra2_acc_producer_state,
|
|
intra2_q_consumer_state,
|
|
x_consumer_state,
|
|
)
|
|
|
|
# Async arrive X/INTRA2_Q/INTRA2_ACC buffer empty/empty/full
|
|
if cutlass.const_expr(self.has_d):
|
|
x_pipeline.consumer_release(
|
|
x_consumer_state, pipeline.PipelineOp.TCGen05Mma
|
|
)
|
|
else:
|
|
x_pipeline.consumer_release(x_consumer_state)
|
|
intra2_q_pipeline.consumer_release(intra2_q_consumer_state)
|
|
intra2_acc_pipeline.producer_commit(intra2_acc_producer_state)
|
|
|
|
# Advance X/INTRA2_Q/INTRA2_ACC state
|
|
x_consumer_state.advance()
|
|
intra2_q_consumer_state.advance()
|
|
intra2_acc_producer_state.advance()
|
|
|
|
# Peek (try_wait) X buffer full for chunk_idx = chunk_idx + 1
|
|
peek_x_full_status = self.conditional_consumer_try_wait(
|
|
x_consumer_state, x_pipeline, C
|
|
)
|
|
|
|
# Advance to next tile
|
|
tile_sched.advance_to_next_work()
|
|
work_tile = tile_sched.get_current_work()
|
|
# END of while work_tile.is_valid_tile
|
|
|
|
# Producer tail for INTRA1_ACC/INTRA2_ACC
|
|
intra1_acc_pipeline.producer_tail(intra1_acc_producer_state)
|
|
intra2_acc_pipeline.producer_tail(intra2_acc_producer_state)
|
|
# END of specialized mma-intra warp
|
|
|
|
# Specialized MMA Inter warp
|
|
elif warp_idx == self.mma_inter_warp_id:
|
|
# Dealloc regs for pre-inter/pre-intra warps
|
|
cute.arch.warpgroup_reg_dealloc(self.num_regs_uniform_warps)
|
|
|
|
# Make shared/tmem fragments for INTER_MMA1 X/B/ACC
|
|
# (MMA, MMA_N, MMA_K, INPUT_STAGE)
|
|
# (MMA, MMA_M, MMA_K, INTERNAL_STAGE)
|
|
# (MMA, MMA_M, MMA_N, INTERNAL_STAGE)
|
|
tCrB, tCrX, tCtAccInter1 = self.mma_partition_ss(
|
|
tiled_mma_inter1,
|
|
self.tile_shape_mnk_inter1,
|
|
smem_bt_internal,
|
|
smem_x,
|
|
tmem_ptr_base + self.tmem_inter1_acc_offset,
|
|
self.internal_stages,
|
|
)
|
|
|
|
# Make shared/tmem fragments for INTER_MMA2 C/P/ACC
|
|
# (MMA, MMA_M, MMA_K, INPUT_STAGE)
|
|
# (MMA, MMA_N, MMA_K, INTERNAL_STAGE)
|
|
# (MMA, MMA_M, MMA_N, INTERNAL_STAGE)
|
|
tCrC, tCrP, tCtAccInter2 = self.mma_partition_ss(
|
|
tiled_mma_inter2,
|
|
self.tile_shape_mnk_inter2,
|
|
smem_c,
|
|
smem_p,
|
|
tmem_ptr_base + self.tmem_inter2_acc_offset,
|
|
self.internal_stages,
|
|
)
|
|
|
|
# Pipeline X/C/INTER1_B/INTER2_P consumer state
|
|
x_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.input_stages
|
|
)
|
|
c_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.input_stages
|
|
)
|
|
inter1_b_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.internal_stages
|
|
)
|
|
inter2_p_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.internal_stages
|
|
)
|
|
|
|
# Pipeline INTER1_ACC/INTER2_ACC producer state
|
|
inter1_acc_producer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Producer, self.internal_stages
|
|
)
|
|
inter2_acc_producer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Producer, self.internal_stages
|
|
)
|
|
|
|
while work_tile.is_valid_tile:
|
|
# Reset count for pipeline state
|
|
x_consumer_state.reset_count()
|
|
c_consumer_state.reset_count()
|
|
inter1_acc_producer_state.reset_count()
|
|
inter1_b_consumer_state.reset_count()
|
|
inter2_p_consumer_state.reset_count()
|
|
inter2_acc_producer_state.reset_count()
|
|
|
|
# Peek (try_wait) C/INTER2_P/INTER2_ACC buffer full/full/empty status
|
|
peek_c_full_status = self.conditional_consumer_try_wait(
|
|
c_consumer_state, c_pipeline, C
|
|
)
|
|
peek_inter2_p_full_status = self.conditional_consumer_try_wait(
|
|
inter2_p_consumer_state, inter2_p_pipeline, C
|
|
)
|
|
peek_inter2_acc_empty_status = self.conditional_producer_try_acquire(
|
|
inter2_acc_producer_state, inter2_acc_pipeline, C
|
|
)
|
|
|
|
# Batched gemm over C dimension
|
|
for chunk_idx in cutlass.range(C, unroll=1):
|
|
# Conditionally wait for C/INTER2_P/INTER2_ACC buffer full/full/empty
|
|
c_pipeline.consumer_wait(c_consumer_state, peek_c_full_status)
|
|
inter2_p_pipeline.consumer_wait(
|
|
inter2_p_consumer_state, peek_inter2_p_full_status
|
|
)
|
|
inter2_acc_pipeline.producer_acquire(
|
|
inter2_acc_producer_state, peek_inter2_acc_empty_status
|
|
)
|
|
|
|
# INTER MMA2
|
|
tiled_mma_inter2 = self.exec_mma(
|
|
tiled_mma_inter2,
|
|
tCtAccInter2,
|
|
tCrC,
|
|
tCrP,
|
|
inter2_acc_producer_state,
|
|
c_consumer_state,
|
|
inter2_p_consumer_state,
|
|
)
|
|
|
|
# Async arrive C/INTER2_P/INTER2_ACC buffer empty/empty/full
|
|
c_pipeline.consumer_release(c_consumer_state)
|
|
inter2_p_pipeline.consumer_release(inter2_p_consumer_state)
|
|
inter2_acc_pipeline.producer_commit(inter2_acc_producer_state)
|
|
|
|
# Wait for X/INTER1_B/INTER1_ACC buffer full/full/empty
|
|
x_pipeline.consumer_wait(x_consumer_state)
|
|
inter1_b_pipeline.consumer_wait(inter1_b_consumer_state)
|
|
inter1_acc_pipeline.producer_acquire(inter1_acc_producer_state)
|
|
|
|
# INTER MMA1
|
|
tiled_mma_inter1 = self.exec_mma(
|
|
tiled_mma_inter1,
|
|
tCtAccInter1,
|
|
tCrB,
|
|
tCrX,
|
|
inter1_acc_producer_state,
|
|
inter1_b_consumer_state,
|
|
x_consumer_state,
|
|
)
|
|
|
|
# Async arrive X/INTER1_B/INTER1_ACC buffer empty/empty/full
|
|
if cutlass.const_expr(self.has_d):
|
|
x_pipeline.consumer_release(
|
|
x_consumer_state, pipeline.PipelineOp.TCGen05Mma
|
|
)
|
|
else:
|
|
x_pipeline.consumer_release(x_consumer_state)
|
|
inter1_b_pipeline.consumer_release(inter1_b_consumer_state)
|
|
inter1_acc_pipeline.producer_commit(inter1_acc_producer_state)
|
|
|
|
# Advance X/C/INTER1_B/INTER1_ACC/INTER2_P/INTER2_ACC state
|
|
x_consumer_state.advance()
|
|
c_consumer_state.advance()
|
|
inter1_b_consumer_state.advance()
|
|
inter1_acc_producer_state.advance()
|
|
inter2_p_consumer_state.advance()
|
|
inter2_acc_producer_state.advance()
|
|
|
|
# Peek (try_wait) C/INTER2_P/INTER2_ACC buffer full/full/empty for chunk_idx = chunk_idx + 1
|
|
peek_c_full_status = self.conditional_consumer_try_wait(
|
|
c_consumer_state, c_pipeline, C
|
|
)
|
|
peek_inter2_p_full_status = self.conditional_consumer_try_wait(
|
|
inter2_p_consumer_state, inter2_p_pipeline, C
|
|
)
|
|
peek_inter2_acc_empty_status = (
|
|
self.conditional_producer_try_acquire(
|
|
inter2_acc_producer_state, inter2_acc_pipeline, C
|
|
)
|
|
)
|
|
|
|
# Advance to next tile
|
|
tile_sched.advance_to_next_work()
|
|
work_tile = tile_sched.get_current_work()
|
|
|
|
# Producer tail for INTER1_ACC/INTER2_ACC
|
|
inter1_acc_pipeline.producer_tail(inter1_acc_producer_state)
|
|
inter2_acc_pipeline.producer_tail(inter2_acc_producer_state)
|
|
|
|
# Specialized Pre-Inter warp
|
|
elif (
|
|
warp_idx == self.pre_inter_warp_id[0]
|
|
or warp_idx == self.pre_inter_warp_id[1]
|
|
or warp_idx == self.pre_inter_warp_id[2]
|
|
or warp_idx == self.pre_inter_warp_id[3]
|
|
):
|
|
# Alloc regs in pre_inter warps
|
|
cute.arch.warpgroup_reg_alloc(self.num_regs_pre_inter_warps)
|
|
|
|
# Make tiledCopy and partition smem/register tensor for smem load Bt
|
|
# ((S2R_ATOM_V, S2R_REST_V), S2R_M, S2R_N, INPUT_STAGE)
|
|
# ((S2R_ATOM_V, S2R_REST_V), S2R_M, S2R_N)
|
|
tiled_s2r_b, tBsB_s2r, tBrB_s2r = self.pre_inter_smem_load_and_partition_b(
|
|
local_tidx, smem_bt
|
|
)
|
|
|
|
# Partition shared tensor for smem store Bt
|
|
smem_bt_internal_ = cute.make_tensor(
|
|
smem_bt_internal.iterator, smem_bt.layout
|
|
)
|
|
# Make tiledCopy and partition register/smem tensor for smem store Bt
|
|
# ((R2S_ATOM_V, R2S_REST_V), R2S_M, R2S_N)
|
|
# ((R2S_ATOM_V, R2S_REST_V), R2S_M, R2S_N, INTERNAL_STAGE)
|
|
tiled_r2s_b, tBrB_r2s, tBsB_r2s = self.pre_inter_smem_store_and_partition_b(
|
|
local_tidx,
|
|
smem_bt_internal_,
|
|
tiled_s2r_b,
|
|
tBrB_s2r,
|
|
)
|
|
|
|
# (MMA, MMA_M, MMA_K, INPUT_STAGE)
|
|
sDelta = self.pre_inter_make_delta(smem_delta, smem_bt.layout)
|
|
sDeltaA = self.pre_inter_make_delta(smem_cumsum_delta, smem_bt.layout)
|
|
|
|
# Make copy_atom and partition register/smem tensor for smem load/store of Delta/DeltaA
|
|
# ((S2R_ATOM_V, S2R_REST_V), S2R_M, S2R_N, INPUT_STAGE)
|
|
# ((S2R_ATOM_V, S2R_REST_V), S2R_M, S2R_N)
|
|
(
|
|
s2r_atom_delta,
|
|
tBsDelta_s2r,
|
|
tBrDelta_s2r,
|
|
) = self.smem_load_and_partition_delta_d(
|
|
tiled_s2r_b, local_tidx, sDelta, (None, None, None, 0)
|
|
)
|
|
(
|
|
s2r_atom_cumsum,
|
|
tBsDeltaA_s2r,
|
|
tBrDeltaA_s2r,
|
|
) = self.smem_load_and_partition_delta_d(
|
|
tiled_s2r_b, local_tidx, sDeltaA, (None, None, None, 0)
|
|
)
|
|
|
|
# ((R2S_ATOM_V, R2S_REST_V), R2S_M, R2S_N)
|
|
thr_r2s_b = tiled_r2s_b.get_slice(local_tidx)
|
|
tBrDelta_r2s = thr_r2s_b.retile(tBrDelta_s2r)
|
|
tBrDeltaA_r2s = thr_r2s_b.retile(tBrDeltaA_s2r)
|
|
|
|
# Make tmem fragment for INTER1_ACC
|
|
# (MMA, MMA_M, MMA_N, INTERNAL_STAGE)
|
|
tCtAccInter1 = self.mma_partition_c(
|
|
tiled_mma_inter1,
|
|
self.tile_shape_mnk_inter1,
|
|
tmem_ptr_base + self.tmem_inter1_acc_offset,
|
|
self.internal_stages,
|
|
)
|
|
# (M_PER_MMA, N_PER_MMA, INTERNAL_STAGE)
|
|
tInter1 = tCtAccInter1[((None, None), 0, 0, None)]
|
|
|
|
# Make tiledCopy and partition tmem/register tensor for tmem load INTER1_ACC
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N, INTERNAL_STAGE)
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N)
|
|
(
|
|
tiled_t2r_inter1,
|
|
tTR_tP,
|
|
tTR_rP,
|
|
) = self.pre_inter_tmem_load_and_partition_p(local_tidx, tInter1, smem_pt)
|
|
|
|
# Make fragment for register to hold P after post-processing (in acc dtype)
|
|
tState = cute.make_rmem_tensor(tTR_rP.shape, self.acc_dtype)
|
|
|
|
# Make tiledCopy and partition smem/register tensor for smem store INTER2_P
|
|
# ((R2S_ATOM_V, R2S_REST_V), R2S_M, R2S_N)
|
|
# ((R2S_ATOM_V, R2S_REST_V), R2S_M, R2S_N, INTERNAL_STAGE)
|
|
tiled_r2s_p, tRS_rP, tRS_sP = self.smem_store_and_partition_p_y(
|
|
local_tidx, smem_pt, tiled_t2r_inter1
|
|
)
|
|
|
|
# Partition global/shared tensor for P (State)
|
|
# ((ATOM_V, REST_V), INTERNAL_STAGE)
|
|
# ((ATOM_V, REST_V), 1, 1, EH, B)
|
|
bSG_sP, bSG_gP_pre_slice = self.tma_partition_with_shape(
|
|
tma_atom_p,
|
|
tma_tensor_p,
|
|
smem_p_store,
|
|
self.tile_shape_mnk_inter2[1:],
|
|
)
|
|
|
|
# Pipeline B/Delta/INTER1_ACC consumer state
|
|
b_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.input_stages
|
|
)
|
|
deltas_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.input_stages
|
|
)
|
|
inter1_acc_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.internal_stages
|
|
)
|
|
|
|
# Pipeline INTER1_B/INTER2_P producer state
|
|
inter1_b_producer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Producer, self.internal_stages
|
|
)
|
|
inter2_p_producer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Producer, self.internal_stages
|
|
)
|
|
|
|
# Pipeline TMA store P
|
|
tma_p_pipeline = pipeline.PipelineTmaStore.create(
|
|
num_stages=self.internal_stages,
|
|
producer_group=pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id)
|
|
),
|
|
)
|
|
|
|
while work_tile.is_valid_tile:
|
|
b_idx, eh_idx, g_idx = work_tile.tile_idx
|
|
|
|
# Slice global tensor to current tile idx
|
|
# ((ATOM_V, REST_V))
|
|
bSG_gP = bSG_gP_pre_slice[(None, 0, 0, eh_idx, b_idx)]
|
|
|
|
# Reset count for pipeline state
|
|
b_consumer_state.reset_count()
|
|
deltas_consumer_state.reset_count()
|
|
inter1_b_producer_state.reset_count()
|
|
inter1_acc_consumer_state.reset_count()
|
|
inter2_p_producer_state.reset_count()
|
|
|
|
# State (P) init
|
|
tState.fill(0.0)
|
|
|
|
# Peek (try_wait) B/Delta/INTER1_B buffer full/full/empty status
|
|
peek_b_full_status = self.conditional_consumer_try_wait(
|
|
b_consumer_state, b_pipeline, C
|
|
)
|
|
peek_deltas_full_status = self.conditional_consumer_try_wait(
|
|
deltas_consumer_state, deltas_pipeline, C
|
|
)
|
|
peek_wr_inter1_b_empty_status = self.conditional_producer_try_acquire(
|
|
inter1_b_producer_state, inter1_b_pipeline, C
|
|
)
|
|
|
|
# Prefill INTER2_P with 0
|
|
# Wait for INTER2_P buffer empty
|
|
inter2_p_pipeline.producer_acquire(inter2_p_producer_state)
|
|
|
|
tRS_rP.fill(0.0)
|
|
# Copy INTER2_P from register to smem
|
|
inter2_p_coord = (None, None, None, inter2_p_producer_state.index)
|
|
cute.copy(tiled_r2s_p, tRS_rP, tRS_sP[inter2_p_coord])
|
|
|
|
# Fence for shared memory
|
|
cute.arch.fence_proxy(
|
|
cute.arch.ProxyKind.async_shared,
|
|
space=cute.arch.SharedSpace.shared_cta,
|
|
)
|
|
# Async arrive INTER2_P buffer full
|
|
inter2_p_pipeline.producer_commit(inter2_p_producer_state)
|
|
# Advance INTER2_P producer state
|
|
inter2_p_producer_state.advance()
|
|
|
|
# Batched processing over C dimension
|
|
for chunk_idx in cutlass.range(C, unroll=1):
|
|
# Conditionally wait for B/Delta/B_TMEM buffer full/full/empty
|
|
b_pipeline.consumer_wait(b_consumer_state, peek_b_full_status)
|
|
deltas_pipeline.consumer_wait(
|
|
deltas_consumer_state, peek_deltas_full_status
|
|
)
|
|
inter1_b_pipeline.producer_acquire(
|
|
inter1_b_producer_state, peek_wr_inter1_b_empty_status
|
|
)
|
|
|
|
# Load B/Delta/DeltaA/last_column
|
|
b_coord = (None, None, None, b_consumer_state.index)
|
|
delta_coord = (None, None, None, deltas_consumer_state.index)
|
|
cute.copy(tiled_s2r_b, tBsB_s2r[b_coord], tBrB_s2r)
|
|
cute.copy(s2r_atom_delta, tBsDelta_s2r[delta_coord], tBrDelta_s2r)
|
|
cute.copy(
|
|
s2r_atom_cumsum, tBsDeltaA_s2r[delta_coord], tBrDeltaA_s2r
|
|
)
|
|
last_column = smem_cumsum_delta[
|
|
smem_cumsum_delta.shape[0] - 1, deltas_consumer_state.index
|
|
]
|
|
|
|
# Fence for shared memory
|
|
cute.arch.fence_proxy(
|
|
cute.arch.ProxyKind.async_shared,
|
|
space=cute.arch.SharedSpace.shared_cta,
|
|
)
|
|
|
|
# Combine B/Delta/DeltaA/last_column
|
|
tScaledB = self.pre_inter_scale_bt_with_delta(
|
|
tBrB_s2r, tBrDelta_r2s, tBrDeltaA_r2s, last_column
|
|
)
|
|
|
|
# Store scaled B to tBrB_r2s
|
|
for reg_idx in range(cute.size(tBrB_r2s)):
|
|
tBrB_r2s[reg_idx] = tScaledB[reg_idx].to(self.io_dtype)
|
|
|
|
# Store tBrB_r2s to bt_smem_internal
|
|
inter1_b_coord = (None, None, None, inter1_b_producer_state.index)
|
|
cute.copy(tiled_r2s_b, tBrB_r2s, tBsB_r2s[inter1_b_coord])
|
|
|
|
# Fence for shared memory
|
|
cute.arch.fence_proxy(
|
|
cute.arch.ProxyKind.async_shared,
|
|
space=cute.arch.SharedSpace.shared_cta,
|
|
)
|
|
|
|
# Async arrive B/Delta/B_TMEM buffer empty/empty/full
|
|
b_pipeline.consumer_release(
|
|
b_consumer_state, pipeline.PipelineOp.AsyncThread
|
|
)
|
|
deltas_pipeline.consumer_release(deltas_consumer_state)
|
|
inter1_b_pipeline.producer_commit(inter1_b_producer_state)
|
|
|
|
# Wait for INTER1_ACC/INTER2_P buffer full/empty
|
|
inter1_acc_pipeline.consumer_wait(inter1_acc_consumer_state)
|
|
inter2_p_pipeline.producer_acquire(inter2_p_producer_state)
|
|
|
|
# Load INTER1_ACC
|
|
inter1_acc_coord = (
|
|
None,
|
|
None,
|
|
None,
|
|
inter1_acc_consumer_state.index,
|
|
)
|
|
cute.copy(tiled_t2r_inter1, tTR_tP[inter1_acc_coord], tTR_rP)
|
|
|
|
# Fence for TMEM load
|
|
cute.arch.fence_view_async_tmem_load()
|
|
|
|
# Combine INTER1_ACC/last_column/State
|
|
exp_last_column = cute.math.exp(last_column, fastmath=True)
|
|
for reg_idx in range(0, cute.size(tTR_rP), 2):
|
|
(
|
|
tTR_rP[reg_idx],
|
|
tTR_rP[reg_idx + 1],
|
|
) = cute.arch.fma_packed_f32x2(
|
|
(exp_last_column, exp_last_column),
|
|
(tState[reg_idx], tState[reg_idx + 1]),
|
|
(tTR_rP[reg_idx], tTR_rP[reg_idx + 1]),
|
|
)
|
|
|
|
# Store scaled P to tRS_rP
|
|
for reg_idx in range(cute.size(tTR_rP)):
|
|
tRS_rP[reg_idx] = tTR_rP[reg_idx].to(self.io_dtype)
|
|
|
|
# Update old state
|
|
tState.store(tTR_rP.load())
|
|
|
|
# Store INTER2_P
|
|
inter2_p_coord = (None, None, None, inter2_p_producer_state.index)
|
|
cute.copy(tiled_r2s_p, tRS_rP, tRS_sP[inter2_p_coord])
|
|
|
|
# Fence for shared memory
|
|
cute.arch.fence_proxy(
|
|
cute.arch.ProxyKind.async_shared,
|
|
space=cute.arch.SharedSpace.shared_cta,
|
|
)
|
|
|
|
# Async arrive INTER1_ACC/INTER2_P buffer empty/full
|
|
inter1_acc_pipeline.consumer_release(inter1_acc_consumer_state)
|
|
# Last iteration consumer is PRE_INTER warp itself, not MMA_INTER warp
|
|
if inter2_p_producer_state.count < C:
|
|
inter2_p_pipeline.producer_commit(inter2_p_producer_state)
|
|
|
|
# Advance B/Delta/INTER1_B/INTER1_ACC state
|
|
b_consumer_state.advance()
|
|
deltas_consumer_state.advance()
|
|
inter1_b_producer_state.advance()
|
|
inter1_acc_consumer_state.advance()
|
|
# Peek (try_wait) B/Delta/INTER1_B buffer full/full./empty for chunk_idx = chunk_idx + 1
|
|
peek_b_full_status = self.conditional_consumer_try_wait(
|
|
b_consumer_state, b_pipeline, C
|
|
)
|
|
peek_deltas_full_status = self.conditional_consumer_try_wait(
|
|
deltas_consumer_state, deltas_pipeline, C
|
|
)
|
|
peek_wr_inter1_b_empty_status = (
|
|
self.conditional_producer_try_acquire(
|
|
inter1_b_producer_state, inter1_b_pipeline, C
|
|
)
|
|
)
|
|
|
|
# Last iteration producer is PRE_INTER warp itself, not MMA_INTER warp
|
|
if inter2_p_producer_state.count < C:
|
|
# Advance INTER2_P producer state
|
|
inter2_p_producer_state.advance()
|
|
# END of for chunk_idx in cutlass.range(C, unroll=1)
|
|
|
|
# Store last INTER2_P (State) from smem to gmem
|
|
# Wait for all previous stores to smem to be done
|
|
cute.arch.fence_proxy(
|
|
cute.arch.ProxyKind.async_shared,
|
|
space=cute.arch.SharedSpace.shared_cta,
|
|
)
|
|
self.pre_inter_sync_barrier.arrive_and_wait()
|
|
|
|
if local_warp_idx == 0:
|
|
# TMA store P
|
|
cute.copy(
|
|
tma_atom_p,
|
|
bSG_sP[(None, inter2_p_producer_state.index)],
|
|
bSG_gP,
|
|
)
|
|
# Wait for TMA store done
|
|
tma_p_pipeline.producer_commit()
|
|
tma_p_pipeline.producer_acquire()
|
|
|
|
self.pre_inter_sync_barrier.arrive_and_wait()
|
|
tma_p_pipeline.producer_tail()
|
|
|
|
# Advance to next tile
|
|
tile_sched.advance_to_next_work()
|
|
work_tile = tile_sched.get_current_work()
|
|
# END of while work_tile.is_valid_tile
|
|
|
|
# Producer tail for INTER1_B/INTER2_P/TMA store P
|
|
inter1_b_pipeline.producer_tail(inter1_b_producer_state)
|
|
inter2_p_pipeline.producer_tail(inter2_p_producer_state)
|
|
# END of specialized pre-inter warp
|
|
|
|
# Specialized Pre-Intra warp
|
|
elif (
|
|
warp_idx == self.pre_intra_warp_id[0]
|
|
or warp_idx == self.pre_intra_warp_id[1]
|
|
or warp_idx == self.pre_intra_warp_id[2]
|
|
or warp_idx == self.pre_intra_warp_id[3]
|
|
):
|
|
# Alloc regs in pre_inter warps
|
|
cute.arch.warpgroup_reg_alloc(self.num_regs_pre_intra_warps)
|
|
|
|
# Make tmem fragment for INTRA1_ACC
|
|
# (MMA, MMA_M, MMA_N, INTRA1_ACC_STAGE)
|
|
tCtAccIntra1 = self.mma_partition_c(
|
|
tiled_mma_intra1,
|
|
self.tile_shape_mnk_intra1,
|
|
tmem_ptr_base + self.tmem_intra1_acc_offset,
|
|
self.intra1_acc_stages,
|
|
)
|
|
# (M_PER_MMA, N_PER_MMA, INTRA1_ACC_STAGE)
|
|
tIntra1 = tCtAccIntra1[((None, None), 0, 0, None)]
|
|
|
|
# Make tiledCopy and partition tmem/register tensor for tensor memory load INTRA1_ACC
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N, INTERNAL_STAGE)
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N)
|
|
tiled_t2r_intra1, tTR_tQ, tTR_rQ = self.pre_intra_tmem_load_and_partition_q(
|
|
tIntra1, local_tidx
|
|
)
|
|
|
|
# Broadcast delta/delta_cumsum smem tensor from LxINPUT_STAGE to LxLxINPUT_STAGE
|
|
sDeltaA_Row = self.pre_intra_make_delta(smem_cumsum_delta, 0)
|
|
sDeltaA_Col = self.pre_intra_make_delta(smem_cumsum_delta, 1)
|
|
sDelta = self.pre_intra_make_delta(smem_delta, 0)
|
|
|
|
# Make tiledCopy and partition smem/register tensor for smem memory load delta/delta_cumsum
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N, INPUT_STAGE)
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N)
|
|
(
|
|
s2r_atom_cumsum,
|
|
tQsDeltaA_Row,
|
|
tQrDeltaA_Row,
|
|
) = self.smem_load_and_partition_delta_d(
|
|
tiled_t2r_intra1, local_tidx, sDeltaA_Row, (None, None, None, 0)
|
|
)
|
|
(
|
|
s2r_atom_cumsum,
|
|
tQsDeltaA_Col,
|
|
tQrDeltaA_Col,
|
|
) = self.smem_load_and_partition_delta_d(
|
|
tiled_t2r_intra1, local_tidx, sDeltaA_Col, (None, None, None, 0)
|
|
)
|
|
(
|
|
s2r_atom_delta,
|
|
tQsDelta,
|
|
tQrDelta,
|
|
) = self.smem_load_and_partition_delta_d(
|
|
tiled_t2r_intra1, local_tidx, sDelta, (None, None, None, 0)
|
|
)
|
|
|
|
# Make and partition coord tensor for delta_cumsum load
|
|
# (L, L)
|
|
coord_tensor = cute.make_identity_tensor(
|
|
cute.dice(self.tile_shape_mnk_intra1, (1, 1, None))
|
|
)
|
|
thr_t2r_intra1 = tiled_t2r_intra1.get_slice(local_tidx)
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N)
|
|
tCoord = thr_t2r_intra1.partition_D(coord_tensor)
|
|
|
|
# Make tmem tensor for INTRA2_Q
|
|
# (MMA, MMA_M, MMA_K, INTERNAL_STAGE)
|
|
tCrQ = self.mma_partition_a_tmem(
|
|
tiled_mma_intra2,
|
|
q_tmem_layout,
|
|
tmem_ptr_base + self.tmem_intra2_q_offset,
|
|
)
|
|
|
|
# Make tiledCopy and partition tmem/register tensor for tensor memory store INTRA2_Q
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N, ...)
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N, ..., INTERNAL_STAGE)
|
|
tiled_r2t_q, tRT_rQ, tRT_tQ = self.pre_intra_tmem_store_and_partition_q(
|
|
local_tidx, tCrQ
|
|
)
|
|
|
|
# Pipeline DELTA/INTRA1_ACC consumer state
|
|
deltas_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.input_stages
|
|
)
|
|
intra1_acc_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.intra1_acc_stages
|
|
)
|
|
# Pipeline INTRA2_Q producer state
|
|
intra2_q_producer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Producer, self.internal_stages
|
|
)
|
|
|
|
while work_tile.is_valid_tile:
|
|
# Reset count for pipeline state
|
|
deltas_consumer_state.reset_count()
|
|
intra1_acc_consumer_state.reset_count()
|
|
intra2_q_producer_state.reset_count()
|
|
|
|
# Peek (try_wait) DELTA/INTRA1_ACC buffer full
|
|
peek_deltas_full_status = self.conditional_consumer_try_wait(
|
|
deltas_consumer_state, deltas_pipeline, C
|
|
)
|
|
peek_rd_intra1_acc_full_status = self.conditional_consumer_try_wait(
|
|
intra1_acc_consumer_state, intra1_acc_pipeline, C
|
|
)
|
|
|
|
# Batched processing over C dimension
|
|
for chunk_idx in cutlass.range(C, unroll=1):
|
|
# Conditionally wait for Delta/INTRA1_ACC buffer full
|
|
deltas_pipeline.consumer_wait(
|
|
deltas_consumer_state, peek_deltas_full_status
|
|
)
|
|
intra1_acc_pipeline.consumer_wait(
|
|
intra1_acc_consumer_state, peek_rd_intra1_acc_full_status
|
|
)
|
|
|
|
# Load Q from tmem
|
|
intra1_coord = (None, None, None, intra1_acc_consumer_state.index)
|
|
cute.copy(tiled_t2r_intra1, tTR_tQ[intra1_coord], tTR_rQ)
|
|
cute.arch.fence_view_async_tmem_load()
|
|
|
|
# Load tQsDeltaA_Row/tQsDeltaA_Col/tQsDelta from smem
|
|
delta_coord = (None, None, None, deltas_consumer_state.index)
|
|
cute.copy(
|
|
s2r_atom_cumsum, tQsDeltaA_Row[delta_coord], tQrDeltaA_Row
|
|
)
|
|
cute.copy(
|
|
s2r_atom_cumsum, tQsDeltaA_Col[delta_coord], tQrDeltaA_Col
|
|
)
|
|
cute.copy(s2r_atom_delta, tQsDelta[delta_coord], tQrDelta)
|
|
|
|
# SegSum
|
|
tRT_rQ = self.pre_intra_segsum(
|
|
tTR_rQ, tQrDeltaA_Row, tQrDeltaA_Col, tQrDelta, tCoord, tRT_rQ
|
|
)
|
|
|
|
# Wait for INTRA2_Q buffer empty
|
|
# Delay producer_acquire to right before data store
|
|
intra2_q_pipeline.producer_acquire(intra2_q_producer_state)
|
|
|
|
# Store Q from reg to tmem
|
|
q_coord = (None, None, None, None, intra2_q_producer_state.index)
|
|
cute.copy(tiled_r2t_q, tRT_rQ, tRT_tQ[q_coord])
|
|
|
|
# Async arrive Delta/INTRA1_ACC buffer empty
|
|
intra1_acc_pipeline.consumer_release(intra1_acc_consumer_state)
|
|
deltas_pipeline.consumer_release(deltas_consumer_state)
|
|
|
|
cute.arch.fence_view_async_tmem_store()
|
|
|
|
# Async arrive INTRA2_Q buffer full
|
|
intra2_q_pipeline.producer_commit(intra2_q_producer_state)
|
|
|
|
# Advance deltas/intra1_acc/intra2_q states
|
|
deltas_consumer_state.advance()
|
|
intra1_acc_consumer_state.advance()
|
|
intra2_q_producer_state.advance()
|
|
|
|
# Peek (try_wait) Delta/INTRA1_ACC buffer full for chunk_idx = chunk_idx + 1
|
|
peek_deltas_full_status = self.conditional_consumer_try_wait(
|
|
deltas_consumer_state, deltas_pipeline, C
|
|
)
|
|
peek_rd_intra1_acc_full_status = self.conditional_consumer_try_wait(
|
|
intra1_acc_consumer_state, intra1_acc_pipeline, C
|
|
)
|
|
# END of for chunk_idx in cutlass.range(C, unroll=1)
|
|
|
|
# Advance to next tile
|
|
tile_sched.advance_to_next_work()
|
|
work_tile = tile_sched.get_current_work()
|
|
# END of while work_tile.is_valid_tile
|
|
|
|
# Producer tail for INTRA2_Q
|
|
intra2_q_pipeline.producer_tail(intra2_q_producer_state)
|
|
# END of specialized pre-intra warp
|
|
|
|
# Specialized Epilogue warp
|
|
else:
|
|
# Dealloc regs for pre-inter/pre-intra warps
|
|
cute.arch.warpgroup_reg_dealloc(self.num_regs_epilogue_warps)
|
|
|
|
# (L, D, INPUT_STAGE)
|
|
sDeltaA = self.epilog_make_delta(smem_cumsum_delta)
|
|
|
|
# Make tmem tensor for INTRA2_ACC/INTER2_ACC
|
|
# (MMA, MMA_M, MMA_K, INTERNAL_STAGE)
|
|
tCtAccIntra2 = self.mma_partition_c(
|
|
tiled_mma_intra2,
|
|
self.tile_shape_mnk_intra2,
|
|
tmem_ptr_base + self.tmem_intra2_acc_offset,
|
|
self.internal_stages,
|
|
)
|
|
# (M_PER_MMA, N_PER_MMA, INTERNAL_STAGE)
|
|
tIntra2 = tCtAccIntra2[((None, None), 0, 0, None)]
|
|
# (MMA, MMA_M, MMA_K, INTERNAL_STAGE)
|
|
tCtAccInter2 = self.mma_partition_c(
|
|
tiled_mma_inter2,
|
|
self.tile_shape_mnk_inter2,
|
|
tmem_ptr_base + self.tmem_inter2_acc_offset,
|
|
self.internal_stages,
|
|
)
|
|
# (M_PER_MMA, N_PER_MMA, INTERNAL_STAGE)
|
|
tInter2 = tCtAccInter2[((None, None), 0, 0, None)]
|
|
|
|
# Subtiling INTRA2_ACC/INTER2_ACC/Delta/Y
|
|
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, INTERNAL_STAGE)
|
|
tIntra_epi = cute.flat_divide(tIntra2, epi_tile)
|
|
tInter_epi = cute.flat_divide(tInter2, epi_tile)
|
|
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, INPUT_STAGE)
|
|
sDeltaA_epi = cute.flat_divide(sDeltaA, epi_tile)
|
|
|
|
# Make tiled copy and partition tmem/reg tensor w.r.t tensor memory load
|
|
# ((T2R_ATOM_V, T2R_REST_V), REST_M, REST_N, EPI_M, EPI_N, INTERNAL_STAGE)
|
|
# ((T2R_ATOM_V, T2R_REST_V), REST_M, REST_N)
|
|
(
|
|
tiled_t2r_intra2,
|
|
tTR_tIntra,
|
|
tTR_rIntra,
|
|
) = self.epilog_tmem_load_and_partition_acc(local_tidx, tIntra_epi, smem_y)
|
|
(
|
|
tiled_t2r_inter2,
|
|
tTR_tInter2,
|
|
tTR_rInter,
|
|
) = self.epilog_tmem_load_and_partition_acc(local_tidx, tInter_epi, smem_y)
|
|
|
|
# Make tiled copy and partition smem/reg tensor w.r.t smem load Delta
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N, EPI_M, EPI_N, INPUT_STAGE)
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N)
|
|
(
|
|
s2r_atom_delta,
|
|
tTR_sDeltaA,
|
|
tTR_rDeltaA,
|
|
) = self.smem_load_and_partition_delta_d(
|
|
tiled_t2r_inter2, local_tidx, sDeltaA_epi, (None, None, None, 0, 0, 0)
|
|
)
|
|
|
|
# Make tiled copy and Partition smem/register tensor w.r.t smem store Y
|
|
# ((R2S_ATOM_V, R2S_REST_V), REST_M, REST_N, OUTPUT_STAGE)
|
|
# ((R2S_ATOM_V, R2S_REST_V), REST_M, REST_N)
|
|
tiled_r2s_y, tRS_rY, tRS_sY = self.smem_store_and_partition_p_y(
|
|
local_tidx, smem_y, tiled_t2r_inter2
|
|
)
|
|
|
|
tRS_rCompute = cute.make_rmem_tensor(tRS_rY.shape, self.acc_dtype)
|
|
|
|
tiled_s2r_x = None
|
|
tSR_sX = None
|
|
tSR_rX = None
|
|
if cutlass.const_expr(self.has_d):
|
|
# Make TiledCopy/smem/register tensor for smem load X
|
|
# (R2S_ATOM, R2S_M, R2S_N, EPI_M, EPI_N, INPUT_STAGES)
|
|
# (R2S_ATOM, R2S_M, R2S_N)
|
|
tiled_s2r_x, tSR_sX, tSR_rX = self.epilog_smem_load_and_partition_x(
|
|
tiled_t2r_inter2, local_tidx, smem_xt, epi_tile
|
|
)
|
|
|
|
tRS_sD = None
|
|
tRS_rD = None
|
|
s2r_atom_d = None
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
# (L, D, INPUT_STAGE)
|
|
sD = self.epilog_make_d(smem_d)
|
|
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, INPUT_STAGE)
|
|
tD_sepi = cute.flat_divide(sD, epi_tile)
|
|
|
|
# Make tiled copy and partition smem/reg tensor w.r.t smem load D
|
|
# ((T2R_ATOM_V, T2R_REST_V), REST_M, REST_N, EPI_M, EPI_N, INPUT_STAGE)
|
|
# ((T2R_ATOM_V, T2R_REST_V), REST_M, REST_N)
|
|
s2r_atom_d, tRS_sD, tRS_rD = self.smem_load_and_partition_delta_d(
|
|
tiled_t2r_inter2, local_tidx, tD_sepi, (None, None, None, 0, 0, 0)
|
|
)
|
|
|
|
elif cutlass.const_expr(self.has_d):
|
|
tRS_rD = cutlass.Float32(0.0).to(self.io_dtype)
|
|
|
|
# Partition global/shared tensor for TMA store Y
|
|
# ((ATOM_V, REST_V), INPUT_STAGE)
|
|
# ((ATOM_V, REST_V), EPI_M, EPI_N, 1, 1, C, EH, B)
|
|
bSG_sY, bSG_gY_pre_slice = self.epilog_tma_partition_y(
|
|
tma_tensor_y, tma_atom_y, smem_y, epi_tile
|
|
)
|
|
|
|
# Make TMA store pipeline Y
|
|
tma_y_pipeline = pipeline.PipelineTmaStore.create(
|
|
num_stages=self.output_stages,
|
|
producer_group=pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, 32 * len(self.epilog_warp_id)
|
|
),
|
|
)
|
|
|
|
# Make consumer pipeline states for Delta/INTRA2_ACC/INTER2_ACC/X/D buffer
|
|
deltas_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.input_stages
|
|
)
|
|
intra2_acc_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.internal_stages
|
|
)
|
|
inter2_acc_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.internal_stages
|
|
)
|
|
x_consumer_state = None
|
|
if cutlass.const_expr(self.has_d):
|
|
x_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.input_stages
|
|
)
|
|
d_consumer_state = None
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
d_consumer_state = pipeline.make_pipeline_state(
|
|
pipeline.PipelineUserType.Consumer, self.input_stages
|
|
)
|
|
|
|
while work_tile.is_valid_tile:
|
|
b_idx, eh_idx, g_idx = work_tile.tile_idx
|
|
|
|
# Slice global tensor to current tile idx
|
|
# ((ATOM_V, REST_V), EPI_M, EPI_N, C)
|
|
bSG_gY = bSG_gY_pre_slice[(None, None, None, 0, 0, None, eh_idx, b_idx)]
|
|
if cutlass.const_expr(self.has_d and not self.d_has_hdim):
|
|
tRS_rD = tma_tensor_d[0, eh_idx]
|
|
|
|
# Reset count for pipeline state
|
|
deltas_consumer_state.reset_count()
|
|
intra2_acc_consumer_state.reset_count()
|
|
inter2_acc_consumer_state.reset_count()
|
|
if cutlass.const_expr(self.has_d):
|
|
x_consumer_state.reset_count()
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
d_consumer_state.reset_count()
|
|
|
|
# Peek Delta/INTRA2_ACC/INTER2_ACC buffer status
|
|
peek_deltas_full_status = self.conditional_consumer_try_wait(
|
|
deltas_consumer_state, deltas_pipeline, C
|
|
)
|
|
peek_rd_intra2_acc_full_status = self.conditional_consumer_try_wait(
|
|
intra2_acc_consumer_state, intra2_acc_pipeline, C
|
|
)
|
|
peek_rd_inter2_acc_full_status = self.conditional_consumer_try_wait(
|
|
inter2_acc_consumer_state, inter2_acc_pipeline, C
|
|
)
|
|
peek_rd_x_full_status = None
|
|
if cutlass.const_expr(self.has_d):
|
|
peek_rd_x_full_status = self.conditional_consumer_try_wait(
|
|
x_consumer_state, x_pipeline, C
|
|
)
|
|
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
d_pipeline.consumer_wait(d_consumer_state)
|
|
|
|
# Batched processing over C dimension
|
|
for chunk_idx in cutlass.range(C, unroll=1):
|
|
# Conditionally wait for Delta/INTRA2_ACC/INTER2_ACC/X buffer full
|
|
deltas_pipeline.consumer_wait(
|
|
deltas_consumer_state, peek_deltas_full_status
|
|
)
|
|
intra2_acc_pipeline.consumer_wait(
|
|
intra2_acc_consumer_state, peek_rd_intra2_acc_full_status
|
|
)
|
|
inter2_acc_pipeline.consumer_wait(
|
|
inter2_acc_consumer_state, peek_rd_inter2_acc_full_status
|
|
)
|
|
if cutlass.const_expr(self.has_d):
|
|
x_pipeline.consumer_wait(
|
|
x_consumer_state, peek_rd_x_full_status
|
|
)
|
|
# Loop over EPI_M and EPI_N subtiles
|
|
for epi_n in range(cute.size(tTR_tIntra, mode=[4])):
|
|
for epi_m in range(cute.size(tTR_tIntra, mode=[3])):
|
|
epi_iter_cnt = (
|
|
epi_n * cute.size(tTR_tIntra, mode=[3]) + epi_m
|
|
)
|
|
epi_buffer_idx = epi_iter_cnt % self.output_stages
|
|
|
|
# Load INTRA2_ACC/INTER2_ACC from tmem
|
|
subtile_coord = (
|
|
None,
|
|
None,
|
|
None,
|
|
epi_m,
|
|
epi_n,
|
|
)
|
|
intra2_coord = subtile_coord + (
|
|
intra2_acc_consumer_state.index,
|
|
)
|
|
cute.copy(
|
|
tiled_t2r_intra2,
|
|
tTR_tIntra[intra2_coord],
|
|
tTR_rIntra,
|
|
)
|
|
inter2_coord = subtile_coord + (
|
|
inter2_acc_consumer_state.index,
|
|
)
|
|
cute.copy(
|
|
tiled_t2r_inter2,
|
|
tTR_tInter2[inter2_coord],
|
|
tTR_rInter,
|
|
)
|
|
# Fence for T2R load
|
|
cute.arch.fence_view_async_tmem_load()
|
|
|
|
# Load Delta from smem
|
|
delta_coord = subtile_coord + (deltas_consumer_state.index,)
|
|
cute.copy(
|
|
s2r_atom_delta, tTR_sDeltaA[delta_coord], tTR_rDeltaA
|
|
)
|
|
|
|
# Load X from smem
|
|
if cutlass.const_expr(self.has_d):
|
|
x_coord = subtile_coord + (x_consumer_state.index,)
|
|
cute.copy(tiled_s2r_x, tSR_sX[x_coord], tSR_rX)
|
|
|
|
# Load D from smem
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
# Load vector D from smem (d_has_hdim = True)
|
|
d_coord = subtile_coord + (d_consumer_state.index,)
|
|
cute.copy(s2r_atom_d, tRS_sD[d_coord], tRS_rD)
|
|
|
|
# Combine INTRA2_ACC/INTER2_ACC/Delta/X/D
|
|
for reg_idx in range(0, cute.size(tRS_rCompute), 2):
|
|
(
|
|
tRS_rCompute[reg_idx],
|
|
tRS_rCompute[reg_idx + 1],
|
|
) = cute.arch.fma_packed_f32x2(
|
|
(tTR_rInter[reg_idx], tTR_rInter[reg_idx + 1]),
|
|
(
|
|
cute.math.exp(
|
|
tTR_rDeltaA[reg_idx], fastmath=True
|
|
),
|
|
cute.math.exp(
|
|
tTR_rDeltaA[reg_idx + 1], fastmath=True
|
|
),
|
|
),
|
|
(tTR_rIntra[reg_idx], tTR_rIntra[reg_idx + 1]),
|
|
)
|
|
# Fuse Y += X * D
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
(
|
|
tRS_rCompute[reg_idx],
|
|
tRS_rCompute[reg_idx + 1],
|
|
) = cute.arch.fma_packed_f32x2(
|
|
(
|
|
tRS_rD[reg_idx].to(self.acc_dtype),
|
|
tRS_rD[reg_idx + 1].to(self.acc_dtype),
|
|
),
|
|
(
|
|
tSR_rX[reg_idx].to(self.acc_dtype),
|
|
tSR_rX[reg_idx + 1].to(self.acc_dtype),
|
|
),
|
|
(
|
|
tRS_rCompute[reg_idx],
|
|
tRS_rCompute[reg_idx + 1],
|
|
),
|
|
)
|
|
elif cutlass.const_expr(self.has_d):
|
|
(
|
|
tRS_rCompute[reg_idx],
|
|
tRS_rCompute[reg_idx + 1],
|
|
) = cute.arch.fma_packed_f32x2(
|
|
(
|
|
tRS_rD.to(self.acc_dtype),
|
|
tRS_rD.to(self.acc_dtype),
|
|
),
|
|
(
|
|
tSR_rX[reg_idx].to(self.acc_dtype),
|
|
tSR_rX[reg_idx + 1].to(self.acc_dtype),
|
|
),
|
|
(
|
|
tRS_rCompute[reg_idx],
|
|
tRS_rCompute[reg_idx + 1],
|
|
),
|
|
)
|
|
|
|
tRS_rY.store(tRS_rCompute.load().to(self.io_dtype))
|
|
|
|
# Store Y to smem
|
|
cute.copy(
|
|
tiled_r2s_y,
|
|
tRS_rY,
|
|
tRS_sY[None, None, None, epi_buffer_idx],
|
|
)
|
|
|
|
# Fence for R2S store
|
|
cute.arch.fence_proxy(
|
|
cute.arch.ProxyKind.async_shared,
|
|
space=cute.arch.SharedSpace.shared_cta,
|
|
)
|
|
# Sync before TMA store
|
|
self.epilog_sync_barrier.arrive_and_wait()
|
|
|
|
# Async arrive Delta/INTRA2_ACC/INTER2_ACC buffer empty
|
|
if (
|
|
epi_iter_cnt
|
|
== cute.size(tTR_tIntra, mode=[4])
|
|
* cute.size(tTR_tIntra, mode=[3])
|
|
- 1
|
|
):
|
|
deltas_pipeline.consumer_release(deltas_consumer_state)
|
|
intra2_acc_pipeline.consumer_release(
|
|
intra2_acc_consumer_state
|
|
)
|
|
inter2_acc_pipeline.consumer_release(
|
|
inter2_acc_consumer_state
|
|
)
|
|
if cutlass.const_expr(self.has_d):
|
|
x_pipeline.consumer_release(
|
|
x_consumer_state,
|
|
pipeline.PipelineOp.AsyncThread,
|
|
)
|
|
|
|
# TMA store Y to global memory
|
|
if local_warp_idx == 0:
|
|
cute.copy(
|
|
tma_atom_y,
|
|
bSG_sY[None, epi_buffer_idx],
|
|
bSG_gY[None, epi_m, epi_n, chunk_idx],
|
|
)
|
|
|
|
# Commit TMA store
|
|
tma_y_pipeline.producer_commit()
|
|
# Wait for TMA store
|
|
tma_y_pipeline.producer_acquire()
|
|
# Sync before smem store
|
|
self.epilog_sync_barrier.arrive_and_wait()
|
|
|
|
# Advance deltas/intra2_acc/inter2_acc consumer states
|
|
deltas_consumer_state.advance()
|
|
intra2_acc_consumer_state.advance()
|
|
inter2_acc_consumer_state.advance()
|
|
|
|
# Peek (try_wait) Delta/INTRA2_ACC/INTER2_ACC buffer full for chunk_idx = chunk_idx + 1
|
|
peek_deltas_full_status = self.conditional_consumer_try_wait(
|
|
deltas_consumer_state, deltas_pipeline, C
|
|
)
|
|
peek_rd_intra2_acc_full_status = self.conditional_consumer_try_wait(
|
|
intra2_acc_consumer_state, intra2_acc_pipeline, C
|
|
)
|
|
peek_rd_inter2_acc_full_status = self.conditional_consumer_try_wait(
|
|
inter2_acc_consumer_state, inter2_acc_pipeline, C
|
|
)
|
|
|
|
if cutlass.const_expr(self.has_d):
|
|
# Advance x consumer states
|
|
x_consumer_state.advance()
|
|
# Peek (try_wait) X buffer full for chunk_idx = chunk_idx + 1
|
|
peek_rd_x_full_status = self.conditional_consumer_try_wait(
|
|
x_consumer_state, x_pipeline, C
|
|
)
|
|
|
|
if cutlass.const_expr(self.d_has_hdim):
|
|
d_pipeline.consumer_release(d_consumer_state)
|
|
d_consumer_state.advance()
|
|
|
|
# Advance to next tile
|
|
tile_sched.advance_to_next_work()
|
|
work_tile = tile_sched.get_current_work()
|
|
|
|
# Producer tail for TMA store Y
|
|
tma_y_pipeline.producer_tail()
|
|
|
|
# Release tensor memory allocation lock
|
|
tmem.relinquish_alloc_permit()
|
|
# Sync before deallocating tmem
|
|
self.tmem_dealloc_sync_barrier.arrive_and_wait()
|
|
# Dealloc tmem buffer
|
|
tmem.free(tmem_ptr_base)
|
|
|
|
return
|
|
|
|
@staticmethod
|
|
def _compute_stages(smem_capacity):
|
|
return 2, 2, 1, 2 # input, output, internal, intra1_acc
|
|
|
|
@staticmethod
|
|
def _compute_grid(y, b, max_active_clusters):
|
|
B = cute.size(y, mode=[4])
|
|
EH = cute.size(y, mode=[3])
|
|
G = cute.size(b, mode=[3])
|
|
NGROUP_RATIO = EH // G
|
|
num_blocks = B * EH
|
|
|
|
tile_sched_params = Mamba2SSDTileSchedulerParams(num_blocks, EH, NGROUP_RATIO)
|
|
grid = Mamba2SSDTileScheduler.get_grid_shape(
|
|
tile_sched_params, max_active_clusters
|
|
)
|
|
return tile_sched_params, grid
|
|
|
|
@staticmethod
|
|
def _plan_tmem_offsets(
|
|
tiled_mma_intra1,
|
|
tile_shape_mnk_intra1,
|
|
tiled_mma_intra2,
|
|
tile_shape_mnk_intra2,
|
|
tiled_mma_inter1,
|
|
tile_shape_mnk_inter1,
|
|
tiled_mma_inter2,
|
|
tile_shape_mnk_inter2,
|
|
acc_stages,
|
|
intra2_a_tmem_layout,
|
|
a_dtype,
|
|
internal_stages,
|
|
intra1_acc_stages,
|
|
):
|
|
SM100_TMEM_CAPACITY_COLUMNS = 512
|
|
BITS_PER_TMEM_COL = 32
|
|
# (MMA, MMA_M, MMA_N)
|
|
acc_shape_intra1 = tiled_mma_intra1.partition_shape_C(tile_shape_mnk_intra1[:2])
|
|
# (MMA, MMA_M, MMA_N)
|
|
tCtAccIntra1_fake = tiled_mma_intra1.make_fragment_C(
|
|
cute.append(acc_shape_intra1, intra1_acc_stages)
|
|
)
|
|
num_intra1_acc_cols = tcgen05.find_tmem_tensor_col_offset(tCtAccIntra1_fake)
|
|
assert tile_shape_mnk_intra1[1] * intra1_acc_stages == num_intra1_acc_cols
|
|
# (MMA, MMA_N, MMA_K, STAGE)
|
|
tCrQ_fake = tiled_mma_intra2.make_fragment_A(intra2_a_tmem_layout.outer.shape)
|
|
num_intra2_a_cols = tcgen05.find_tmem_tensor_col_offset(tCrQ_fake)
|
|
assert (
|
|
tile_shape_mnk_intra2[2]
|
|
* internal_stages
|
|
* a_dtype.width
|
|
// BITS_PER_TMEM_COL
|
|
== num_intra2_a_cols
|
|
)
|
|
# (MMA, MMA_M, MMA_N)
|
|
acc_shape_intra2 = tiled_mma_intra2.partition_shape_C(tile_shape_mnk_intra2[:2])
|
|
# (MMA, MMA_M, MMA_N)
|
|
tCtAccIntra2_fake = tiled_mma_intra2.make_fragment_C(
|
|
cute.append(acc_shape_intra2, acc_stages)
|
|
)
|
|
num_intra2_acc_cols = tcgen05.find_tmem_tensor_col_offset(tCtAccIntra2_fake)
|
|
assert tile_shape_mnk_intra2[1] * acc_stages == num_intra2_acc_cols
|
|
|
|
# (MMA, MMA_M, MMA_N)
|
|
acc_shape_inter1 = tiled_mma_inter1.partition_shape_C(tile_shape_mnk_inter1[:2])
|
|
# (MMA, MMA_M, MMA_N)
|
|
tCtAccInter1_fake = tiled_mma_inter1.make_fragment_C(
|
|
cute.append(acc_shape_inter1, acc_stages)
|
|
)
|
|
num_inter1_acc_cols = tcgen05.find_tmem_tensor_col_offset(tCtAccInter1_fake)
|
|
assert tile_shape_mnk_inter1[1] * acc_stages == num_inter1_acc_cols
|
|
|
|
# (MMA, MMA_M, MMA_N)
|
|
acc_shape_inter2 = tiled_mma_inter2.partition_shape_C(tile_shape_mnk_inter2[:2])
|
|
# (MMA, MMA_M, MMA_N)
|
|
tCtAccInter2_fake = tiled_mma_inter2.make_fragment_C(
|
|
cute.append(acc_shape_inter2, acc_stages)
|
|
)
|
|
num_inter2_acc_cols = tcgen05.find_tmem_tensor_col_offset(tCtAccInter2_fake)
|
|
assert tile_shape_mnk_inter2[1] * acc_stages == num_inter2_acc_cols
|
|
|
|
tmem_intra1_acc_offset = 0
|
|
tmem_intra2_q_offset = tmem_intra1_acc_offset + num_intra1_acc_cols
|
|
tmem_intra2_acc_offset = tmem_intra2_q_offset + num_intra2_a_cols
|
|
tmem_inter1_acc_offset = tmem_intra2_acc_offset + num_intra2_acc_cols
|
|
tmem_inter2_acc_offset = tmem_inter1_acc_offset + num_inter1_acc_cols
|
|
num_tmem_cols_total_tmp = tmem_inter2_acc_offset + num_inter2_acc_cols
|
|
# Turn num_tmem_cols_total to the nearest power of 2
|
|
num_tmem_cols_total = 1
|
|
while num_tmem_cols_total < num_tmem_cols_total_tmp:
|
|
num_tmem_cols_total *= 2
|
|
assert num_tmem_cols_total <= SM100_TMEM_CAPACITY_COLUMNS
|
|
|
|
return (
|
|
tmem_intra1_acc_offset,
|
|
tmem_intra2_q_offset,
|
|
tmem_intra2_acc_offset,
|
|
tmem_inter1_acc_offset,
|
|
tmem_inter2_acc_offset,
|
|
num_tmem_cols_total,
|
|
)
|
|
|
|
@staticmethod
|
|
def make_tiled_mmas(
|
|
io_dtype,
|
|
acc_dtype,
|
|
cta_group,
|
|
tile_shape_mnk_intra1,
|
|
tile_shape_mnk_intra2,
|
|
tile_shape_mnk_inter1,
|
|
tile_shape_mnk_inter2,
|
|
):
|
|
tiled_mma_intra1 = sm100_utils.make_trivial_tiled_mma(
|
|
io_dtype,
|
|
tcgen05.OperandMajorMode("mn"),
|
|
tcgen05.OperandMajorMode("mn"),
|
|
acc_dtype,
|
|
cta_group,
|
|
tile_shape_mnk_intra1[:2],
|
|
tcgen05.OperandSource.SMEM,
|
|
)
|
|
tiled_mma_intra2 = sm100_utils.make_trivial_tiled_mma(
|
|
io_dtype,
|
|
tcgen05.OperandMajorMode("k"),
|
|
tcgen05.OperandMajorMode("k"),
|
|
acc_dtype,
|
|
cta_group,
|
|
tile_shape_mnk_intra2[:2],
|
|
tcgen05.OperandSource.TMEM,
|
|
)
|
|
tiled_mma_inter1 = sm100_utils.make_trivial_tiled_mma(
|
|
io_dtype,
|
|
tcgen05.OperandMajorMode("k"),
|
|
tcgen05.OperandMajorMode("k"),
|
|
acc_dtype,
|
|
cta_group,
|
|
tile_shape_mnk_inter1[:2],
|
|
tcgen05.OperandSource.SMEM,
|
|
)
|
|
tiled_mma_inter2 = sm100_utils.make_trivial_tiled_mma(
|
|
io_dtype,
|
|
tcgen05.OperandMajorMode("mn"),
|
|
tcgen05.OperandMajorMode("k"),
|
|
acc_dtype,
|
|
cta_group,
|
|
tile_shape_mnk_inter2[:2],
|
|
tcgen05.OperandSource.SMEM,
|
|
)
|
|
return tiled_mma_intra1, tiled_mma_intra2, tiled_mma_inter1, tiled_mma_inter2
|
|
|
|
def make_and_init_x_pipeline(self, x_full_mbar_ptr):
|
|
x_producer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.tma_deltas_x_d_warp_id])
|
|
)
|
|
if not self.has_d:
|
|
x_consumer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread,
|
|
len([self.mma_intra_warp_id, self.mma_inter_warp_id]),
|
|
)
|
|
return pipeline.PipelineTmaUmma.create(
|
|
num_stages=self.input_stages,
|
|
producer_group=x_producer_group,
|
|
consumer_group=x_consumer_group,
|
|
tx_count=self.num_x_load_bytes,
|
|
barrier_storage=x_full_mbar_ptr,
|
|
defer_sync=True,
|
|
)
|
|
else:
|
|
x_consumer_group_umma = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread,
|
|
len([self.mma_intra_warp_id, self.mma_inter_warp_id]),
|
|
)
|
|
x_consumer_group_async = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, 32 * len(self.epilog_warp_id)
|
|
)
|
|
return pipeline.PipelineTmaMultiConsumersAsync.create(
|
|
num_stages=self.input_stages,
|
|
producer_group=x_producer_group,
|
|
consumer_group_umma=x_consumer_group_umma,
|
|
consumer_group_async=x_consumer_group_async,
|
|
tx_count=self.num_x_load_bytes,
|
|
barrier_storage=x_full_mbar_ptr,
|
|
defer_sync=True,
|
|
)
|
|
|
|
def make_and_init_b_pipeline(self, b_full_mbar_ptr):
|
|
b_producer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.tma_b_c_warp_id])
|
|
)
|
|
b_consumer_group_umma = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.mma_intra_warp_id])
|
|
)
|
|
b_consumer_group_async = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id)
|
|
)
|
|
return pipeline.PipelineTmaMultiConsumersAsync.create(
|
|
num_stages=self.input_stages,
|
|
producer_group=b_producer_group,
|
|
consumer_group_umma=b_consumer_group_umma,
|
|
consumer_group_async=b_consumer_group_async,
|
|
tx_count=self.num_b_load_bytes,
|
|
barrier_storage=b_full_mbar_ptr,
|
|
defer_sync=True,
|
|
)
|
|
|
|
def make_and_init_c_pipeline(self, c_full_mbar_ptr):
|
|
c_producer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.tma_b_c_warp_id])
|
|
)
|
|
c_consumer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.mma_intra_warp_id, self.mma_inter_warp_id])
|
|
)
|
|
return pipeline.PipelineTmaUmma.create(
|
|
num_stages=self.input_stages,
|
|
producer_group=c_producer_group,
|
|
consumer_group=c_consumer_group,
|
|
tx_count=self.num_c_load_bytes,
|
|
barrier_storage=c_full_mbar_ptr,
|
|
defer_sync=True,
|
|
)
|
|
|
|
def make_and_init_deltas_pipeline(self, deltas_full_mbar_ptr):
|
|
deltas_producer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.tma_deltas_x_d_warp_id])
|
|
)
|
|
deltas_consumer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread,
|
|
len(
|
|
[*self.pre_inter_warp_id, *self.pre_intra_warp_id, *self.epilog_warp_id]
|
|
),
|
|
)
|
|
|
|
return pipeline.PipelineTmaAsync.create(
|
|
num_stages=self.input_stages,
|
|
producer_group=deltas_producer_group,
|
|
consumer_group=deltas_consumer_group,
|
|
tx_count=self.num_delta_load_bytes + self.num_cumsum_delta_load_bytes,
|
|
barrier_storage=deltas_full_mbar_ptr,
|
|
defer_sync=True,
|
|
)
|
|
|
|
def make_and_init_d_pipeline(self, d_full_mbar_ptr):
|
|
if not self.d_has_hdim:
|
|
return None
|
|
else:
|
|
d_producer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.tma_deltas_x_d_warp_id])
|
|
)
|
|
d_consumer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len(self.epilog_warp_id)
|
|
)
|
|
|
|
return pipeline.PipelineTmaAsync.create(
|
|
num_stages=self.input_stages,
|
|
producer_group=d_producer_group,
|
|
consumer_group=d_consumer_group,
|
|
tx_count=self.num_d_load_bytes,
|
|
barrier_storage=d_full_mbar_ptr,
|
|
defer_sync=True,
|
|
)
|
|
|
|
def make_and_init_intra1_acc_pipeline(self, intra1_acc_full_mbar_ptr):
|
|
intra1_acc_producer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.mma_intra_warp_id])
|
|
)
|
|
intra1_acc_consumer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, 32 * len(self.pre_intra_warp_id)
|
|
)
|
|
return pipeline.PipelineUmmaAsync.create(
|
|
num_stages=self.intra1_acc_stages,
|
|
producer_group=intra1_acc_producer_group,
|
|
consumer_group=intra1_acc_consumer_group,
|
|
barrier_storage=intra1_acc_full_mbar_ptr,
|
|
defer_sync=True,
|
|
)
|
|
|
|
def make_and_init_intra2_q_pipeline(self, intra2_q_full_mbar_ptr):
|
|
intra2_q_producer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, 32 * len(self.pre_intra_warp_id)
|
|
)
|
|
intra2_q_consumer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.mma_intra_warp_id])
|
|
)
|
|
return pipeline.PipelineAsyncUmma.create(
|
|
num_stages=self.internal_stages,
|
|
producer_group=intra2_q_producer_group,
|
|
consumer_group=intra2_q_consumer_group,
|
|
barrier_storage=intra2_q_full_mbar_ptr,
|
|
defer_sync=True,
|
|
)
|
|
|
|
def make_and_init_intra2_acc_pipeline(self, intra2_acc_full_mbar_ptr):
|
|
intra2_acc_producer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.mma_intra_warp_id])
|
|
)
|
|
intra2_acc_consumer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, 32 * len(self.epilog_warp_id)
|
|
)
|
|
return pipeline.PipelineUmmaAsync.create(
|
|
num_stages=self.internal_stages,
|
|
producer_group=intra2_acc_producer_group,
|
|
consumer_group=intra2_acc_consumer_group,
|
|
barrier_storage=intra2_acc_full_mbar_ptr,
|
|
defer_sync=True,
|
|
)
|
|
|
|
def make_and_init_inter1_b_pipeline(self, inter1_b_full_mbar_ptr):
|
|
inter1_b_producer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id)
|
|
)
|
|
inter1_b_consumer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.mma_inter_warp_id])
|
|
)
|
|
return pipeline.PipelineAsyncUmma.create(
|
|
num_stages=self.internal_stages,
|
|
producer_group=inter1_b_producer_group,
|
|
consumer_group=inter1_b_consumer_group,
|
|
barrier_storage=inter1_b_full_mbar_ptr,
|
|
defer_sync=True,
|
|
)
|
|
|
|
def make_and_init_inter1_acc_pipeline(self, inter1_acc_full_mbar_ptr):
|
|
inter1_acc_producer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.mma_inter_warp_id])
|
|
)
|
|
inter1_acc_consumer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id)
|
|
)
|
|
return pipeline.PipelineUmmaAsync.create(
|
|
num_stages=self.internal_stages,
|
|
producer_group=inter1_acc_producer_group,
|
|
consumer_group=inter1_acc_consumer_group,
|
|
barrier_storage=inter1_acc_full_mbar_ptr,
|
|
defer_sync=True,
|
|
)
|
|
|
|
def make_and_init_inter2_p_pipeline(self, inter2_p_full_mbar_ptr):
|
|
inter2_p_producer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, 32 * len(self.pre_inter_warp_id)
|
|
)
|
|
inter2_p_consumer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.mma_inter_warp_id])
|
|
)
|
|
return pipeline.PipelineAsyncUmma.create(
|
|
num_stages=self.internal_stages,
|
|
producer_group=inter2_p_producer_group,
|
|
consumer_group=inter2_p_consumer_group,
|
|
barrier_storage=inter2_p_full_mbar_ptr,
|
|
defer_sync=True,
|
|
)
|
|
|
|
def make_and_init_inter2_acc_pipeline(self, inter2_acc_full_mbar_ptr):
|
|
inter2_acc_producer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, len([self.mma_inter_warp_id])
|
|
)
|
|
inter2_acc_consumer_group = pipeline.CooperativeGroup(
|
|
pipeline.Agent.Thread, 32 * len(self.epilog_warp_id)
|
|
)
|
|
return pipeline.PipelineUmmaAsync.create(
|
|
num_stages=self.internal_stages,
|
|
producer_group=inter2_acc_producer_group,
|
|
consumer_group=inter2_acc_consumer_group,
|
|
barrier_storage=inter2_acc_full_mbar_ptr,
|
|
defer_sync=True,
|
|
)
|
|
|
|
def tma_partition_for_mma_b_operand(
|
|
self,
|
|
tma_atom_x,
|
|
tma_tensor_x,
|
|
smem_x,
|
|
tiled_mma_intra2,
|
|
cluster_layout_vmnk,
|
|
mma_tile_coord_v,
|
|
block_in_cluster_coord_vmnk,
|
|
):
|
|
# Local_tile partition global tensors
|
|
# (D, L, 1, 1, C, EH, B)
|
|
gX = cute.local_tile(
|
|
tma_tensor_x,
|
|
self.tile_shape_mnk_intra2[1:],
|
|
(None, None, None, None, None),
|
|
)
|
|
# Partition global tensor with regard to TiledMMA
|
|
thr_mma_intra2 = tiled_mma_intra2.get_slice(mma_tile_coord_v)
|
|
# (MMA, MMA_N, MMA_K, 1, 1, C, EH, B)
|
|
tCgX = thr_mma_intra2.partition_B(gX)
|
|
|
|
# Partition global/shared tensor for X
|
|
x_cta_layout = cute.make_layout(
|
|
cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
|
|
)
|
|
# ((ATOM_V, REST_V), INPUT_STAGE)
|
|
# ((ATOM_V, REST_V), 1, 1, C, EH, B)
|
|
tXsX, tXgX_pre_slice = cpasync.tma_partition(
|
|
tma_atom_x,
|
|
block_in_cluster_coord_vmnk[2],
|
|
x_cta_layout,
|
|
cute.group_modes(smem_x, 0, 3),
|
|
cute.group_modes(tCgX, 0, 3),
|
|
)
|
|
return tXsX, tXgX_pre_slice
|
|
|
|
def tma_partition_for_mma_a_operand(
|
|
self,
|
|
tma_atom_c,
|
|
tma_tensor_c,
|
|
smem_c,
|
|
tiled_mma_intra1,
|
|
cluster_layout_vmnk,
|
|
mma_tile_coord_v,
|
|
block_in_cluster_coord_vmnk,
|
|
):
|
|
# Local_tile partition global tensors
|
|
# (L, N, 1, 1, C, G, B)
|
|
gC = cute.local_tile(
|
|
tma_tensor_c,
|
|
cute.slice_(self.tile_shape_mnk_intra1, (None, 0, None)),
|
|
(None, None, None, None, None),
|
|
)
|
|
# Partition global tensor with regard to TiledMMA
|
|
thr_mma_intra1 = tiled_mma_intra1.get_slice(mma_tile_coord_v)
|
|
# (MMA, MMA_M/N, MMA_K, 1, 1, C, G, B)
|
|
tCgC = thr_mma_intra1.partition_A(gC)
|
|
|
|
# Partition global/shared tensor for TMA C
|
|
c_cta_layout = cute.make_layout(
|
|
cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape
|
|
)
|
|
# ((ATOM_V, REST_V), INPUT_STAGE)
|
|
# ((ATOM_V, REST_V), 1, 1, C, G, B)
|
|
tCsC, tCgC_pre_slice = cpasync.tma_partition(
|
|
tma_atom_c,
|
|
block_in_cluster_coord_vmnk[1],
|
|
c_cta_layout,
|
|
cute.group_modes(smem_c, 0, 3),
|
|
cute.group_modes(tCgC, 0, 3),
|
|
)
|
|
return tCsC, tCgC_pre_slice
|
|
|
|
def tma_partition_with_shape(
|
|
self, tma_atom_delta, tma_tensor_delta, smem_delta, shape
|
|
):
|
|
# Local_tile partition global tensors
|
|
# (L, 1, C, EH, B)
|
|
gDelta = cute.local_tile(
|
|
tma_tensor_delta,
|
|
shape,
|
|
(None,) * cute.rank(tma_tensor_delta),
|
|
)
|
|
# Partition global/shared tensor for DELTA
|
|
# ((ATOM_V, REST_V), INPUT_STAGE)
|
|
# ((ATOM_V, REST_V), 1, C, EH, B)
|
|
tDeltasDelta, tDeltagDelta_pre_slice = cpasync.tma_partition(
|
|
tma_atom_delta,
|
|
0,
|
|
cute.make_layout(1),
|
|
cute.group_modes(smem_delta, 0, cute.rank(shape)),
|
|
cute.group_modes(gDelta, 0, cute.rank(shape)),
|
|
)
|
|
|
|
return tDeltasDelta, tDeltagDelta_pre_slice
|
|
|
|
def mma_partition_ss(
|
|
self,
|
|
tiled_mma,
|
|
tile_shape_mnk,
|
|
smem_a,
|
|
smem_b,
|
|
tmem_acc_ptr,
|
|
acc_stages,
|
|
):
|
|
# (MMA, MMA_M, MMA_K, INPUT_STAGE)
|
|
tCrA = tiled_mma.make_fragment_A(smem_a)
|
|
# (MMA, MMA_N, MMA_K, INPUT_STAGE)
|
|
tCrB = tiled_mma.make_fragment_B(smem_b)
|
|
# (MMA, MMA_M, MMA_N, ACC_STAGE)
|
|
tCtAcc = self.mma_partition_c(
|
|
tiled_mma, tile_shape_mnk, tmem_acc_ptr, acc_stages
|
|
)
|
|
return tCrA, tCrB, tCtAcc
|
|
|
|
def mma_partition_ts(
|
|
self,
|
|
tiled_mma,
|
|
tile_shape_mnk,
|
|
a_tmem_layout,
|
|
smem_b,
|
|
tmem_a_ptr,
|
|
tmem_acc_ptr,
|
|
acc_stages,
|
|
):
|
|
# (MMA, MMA_M, MMA_K, INTERNAL_STAGE)
|
|
tCrA = self.mma_partition_a_tmem(tiled_mma, a_tmem_layout, tmem_a_ptr)
|
|
# (MMA, MMA_N, MMA_K, INPUT_STAGE)
|
|
tCrB = tiled_mma.make_fragment_B(smem_b)
|
|
# (MMA, MMA_M, MMA_N, INTERNAL_STAGE)
|
|
tCtAcc = self.mma_partition_c(
|
|
tiled_mma, tile_shape_mnk, tmem_acc_ptr, acc_stages
|
|
)
|
|
return tCrA, tCrB, tCtAcc
|
|
|
|
def mma_partition_a_tmem(self, tiled_mma, a_tmem_layout, tmem_a_ptr):
|
|
tCrA_fake = tiled_mma.make_fragment_A(a_tmem_layout.outer.shape)
|
|
tCrA = cute.make_tensor(
|
|
cute.recast_ptr(
|
|
tmem_a_ptr,
|
|
dtype=tCrA_fake.element_type,
|
|
),
|
|
tCrA_fake.layout,
|
|
)
|
|
return tCrA
|
|
|
|
def mma_partition_c(self, tiled_mma, tile_shape_mnk, tmem_acc_ptr, acc_stages):
|
|
acc_shape = tiled_mma.partition_shape_C(tile_shape_mnk[:2])
|
|
tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, acc_stages))
|
|
# (MMA, MMA_M, MMA_N, INTERNAL_STAGE)
|
|
tCtAcc = cute.make_tensor(tmem_acc_ptr, tCtAcc_fake.layout)
|
|
return tCtAcc
|
|
|
|
@cute.jit
|
|
def exec_mma(
|
|
self,
|
|
tiled_mma,
|
|
tCtAcc,
|
|
tCrA,
|
|
tCrB,
|
|
acc_producer_state,
|
|
a_consumer_state,
|
|
b_consumer_state,
|
|
):
|
|
for kphase_idx in cutlass.range(cute.size(tCrB, mode=[2]), unroll_full=True):
|
|
# set accu = 1
|
|
tiled_mma.set(
|
|
tcgen05.Field.ACCUMULATE,
|
|
cutlass.Boolean(kphase_idx != 0),
|
|
)
|
|
cute.gemm(
|
|
tiled_mma,
|
|
tCtAcc[None, None, None, acc_producer_state.index],
|
|
tCrA[None, None, kphase_idx, a_consumer_state.index],
|
|
tCrB[None, None, kphase_idx, b_consumer_state.index],
|
|
tCtAcc[None, None, None, acc_producer_state.index],
|
|
)
|
|
return tiled_mma
|
|
|
|
@cute.jit
|
|
def conditional_consumer_try_wait(self, b_consumer_state, b_pipeline, C):
|
|
peek_b_full_status = cutlass.Boolean(1)
|
|
if b_consumer_state.count < C:
|
|
peek_b_full_status = b_pipeline.consumer_try_wait(b_consumer_state)
|
|
return peek_b_full_status
|
|
|
|
@cute.jit
|
|
def conditional_producer_try_acquire(
|
|
self, intra1_acc_producer_state, intra1_acc_pipeline, C
|
|
):
|
|
peek_wr_intra1_acc_empty_status = cutlass.Boolean(1)
|
|
if intra1_acc_producer_state.count < C:
|
|
peek_wr_intra1_acc_empty_status = intra1_acc_pipeline.producer_try_acquire(
|
|
intra1_acc_producer_state
|
|
)
|
|
return peek_wr_intra1_acc_empty_status
|
|
|
|
def pre_intra_tmem_load_and_partition_q(self, tIntra1, local_tidx):
|
|
copy_atom_t2r_intra1 = cute.make_copy_atom(
|
|
tcgen05.Ld16x256bOp(tcgen05.Repetition(16), tcgen05.Pack.NONE),
|
|
self.acc_dtype,
|
|
)
|
|
# (L, L)
|
|
fake_sQ = cute.make_tensor(
|
|
cute.make_ptr(self.io_dtype, 0, cute.AddressSpace.smem),
|
|
cute.dice(self.tile_shape_mnk_intra1, (1, 1, None)),
|
|
)
|
|
return self.make_tmem_load_and_partition(
|
|
copy_atom_t2r_intra1, tIntra1, (None, None, 0), local_tidx, fake_sQ
|
|
)
|
|
|
|
def pre_intra_make_delta(self, smem_delta, extend_on_row_or_col):
|
|
smem_iterator = smem_delta.iterator
|
|
delta_linear_smem_layout = smem_delta.layout
|
|
# extend L linear layout to LxL
|
|
extend_layout = cute.make_layout(delta_linear_smem_layout.shape[0], stride=0)
|
|
if extend_on_row_or_col == 0:
|
|
# (L, L, INPUT_STAGE):(0, 1, L)
|
|
sDelta = cute.make_tensor(
|
|
smem_iterator,
|
|
cute.prepend(
|
|
delta_linear_smem_layout,
|
|
extend_layout,
|
|
),
|
|
)
|
|
else:
|
|
# (L, L, INPUT_STAGE):(1, 0, L)
|
|
sDelta = cute.make_tensor(
|
|
smem_iterator,
|
|
cute.append(
|
|
cute.append(
|
|
cute.get(delta_linear_smem_layout, mode=[0]),
|
|
extend_layout,
|
|
),
|
|
cute.get(delta_linear_smem_layout, mode=[1]),
|
|
),
|
|
)
|
|
return sDelta
|
|
|
|
def pre_intra_tmem_store_and_partition_q(self, local_tidx, tCrQ):
|
|
dtype = tCrQ.element_type
|
|
# Make tiledCopy for tensor memory store INTRA2_Q
|
|
copy_atom_r2t_q = cute.make_copy_atom(
|
|
tcgen05.St16x128bOp(tcgen05.Repetition(16), tcgen05.Unpack.NONE),
|
|
dtype,
|
|
)
|
|
tiled_r2t_q = tcgen05.make_tmem_copy(copy_atom_r2t_q, tCrQ)
|
|
thr_r2t_q = tiled_r2t_q.get_slice(local_tidx)
|
|
|
|
# Partition tmem/register tensor for tensor memory store INTRA2_Q
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N, ...)
|
|
tRT_rQ = cute.make_rmem_tensor(
|
|
cute.slice_(thr_r2t_q.partition_S(tCrQ).shape, (None, None, None, None, 0)),
|
|
dtype,
|
|
)
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N, ..., INTERNAL_STAGE)
|
|
tRT_tQ = thr_r2t_q.partition_D(tCrQ)
|
|
|
|
return tiled_r2t_q, tRT_rQ, tRT_tQ
|
|
|
|
@cute.jit
|
|
def pre_intra_segsum(
|
|
self, tTR_rQ, tQrDeltaA_Row, tQrDeltaA_Col, tQrDelta, tCoord, tRT_rQ
|
|
):
|
|
# Make tmp acc type fragments
|
|
tCrDeltaA_Row = cute.make_rmem_tensor(tQrDeltaA_Row.shape, self.acc_dtype)
|
|
tCrDeltaA_Col = cute.make_rmem_tensor(tQrDeltaA_Col.shape, self.acc_dtype)
|
|
tCrDelta = cute.make_rmem_tensor(tQrDelta.shape, self.acc_dtype)
|
|
tCompute = cute.make_rmem_tensor(tRT_rQ.shape, self.acc_dtype)
|
|
|
|
# Combine tTR_rQ/tCrDeltaA_Row/tCrDeltaA_Col/tCrDelta
|
|
tCrDeltaA_Row.store(tQrDeltaA_Row.load().to(self.acc_dtype))
|
|
tCrDeltaA_Col.store(tQrDeltaA_Col.load().to(self.acc_dtype))
|
|
tCrDelta.store(tQrDelta.load().to(self.acc_dtype))
|
|
|
|
# SegSum
|
|
# fadd2 + fsel + fmul2/mufu + fmul2
|
|
for subtile_idx in cutlass.range(0, cute.size(tTR_rQ), 2, unroll_full=True):
|
|
(
|
|
tCompute[subtile_idx],
|
|
tCompute[subtile_idx + 1],
|
|
) = cute.arch.add_packed_f32x2(
|
|
(tCrDeltaA_Col[subtile_idx], tCrDeltaA_Col[subtile_idx + 1]),
|
|
(-tCrDeltaA_Row[subtile_idx], -tCrDeltaA_Row[subtile_idx + 1]),
|
|
)
|
|
for subtile_idx in cutlass.range(cute.size(tTR_rQ), unroll_full=True):
|
|
m, n = tCoord[subtile_idx]
|
|
if m < n:
|
|
tCompute[subtile_idx] = cutlass.Float32(-float("inf"))
|
|
LOG2_E = cutlass.Float32(1.4426950408889634)
|
|
for subtile_idx in cutlass.range(0, cute.size(tTR_rQ), 2, unroll_full=True):
|
|
# TODO: use math.exp directly
|
|
tCompute_log2e = cute.arch.mul_packed_f32x2(
|
|
(tCompute[subtile_idx], tCompute[subtile_idx + 1]), (LOG2_E, LOG2_E)
|
|
)
|
|
(
|
|
tCompute[subtile_idx],
|
|
tCompute[subtile_idx + 1],
|
|
) = cute.arch.mul_packed_f32x2(
|
|
(
|
|
cute.math.exp2(tCompute_log2e[0], fastmath=True),
|
|
cute.math.exp2(tCompute_log2e[1], fastmath=True),
|
|
),
|
|
(tCrDelta[subtile_idx], tCrDelta[subtile_idx + 1]),
|
|
)
|
|
(
|
|
tCompute[subtile_idx],
|
|
tCompute[subtile_idx + 1],
|
|
) = cute.arch.mul_packed_f32x2(
|
|
(tCompute[subtile_idx], tCompute[subtile_idx + 1]),
|
|
(tTR_rQ[subtile_idx], tTR_rQ[subtile_idx + 1]),
|
|
)
|
|
|
|
tRT_rQ.store(tCompute.load().to(self.io_dtype))
|
|
return tRT_rQ
|
|
|
|
def pre_inter_smem_load_and_partition_b(self, local_tidx, smem_bt):
|
|
dtype = smem_bt.element_type
|
|
copy_atom_s2r_b = cute.make_copy_atom(
|
|
cute.nvgpu.CopyUniversalOp(),
|
|
dtype,
|
|
num_bits_per_copy=128,
|
|
)
|
|
num_elements_per_thread = 128 // dtype.width
|
|
num_threads_per_row = self.tile_shape_mnk_inter1[2] // num_elements_per_thread
|
|
num_threads_per_col = 128 // num_threads_per_row
|
|
thread_layout = cute.make_layout(
|
|
(num_threads_per_col, num_threads_per_row),
|
|
stride=(num_threads_per_row, 1),
|
|
)
|
|
val_layout = cute.make_layout((1, num_elements_per_thread))
|
|
tiled_s2r_b = cute.make_tiled_copy_tv(
|
|
copy_atom_s2r_b,
|
|
thread_layout,
|
|
val_layout,
|
|
)
|
|
thr_s2r_b = tiled_s2r_b.get_slice(local_tidx)
|
|
|
|
# Partition shared tensor for smem load Bt
|
|
# ((S2R_ATOM_V, S2R_REST_V), S2R_M, S2R_N, INPUT_STAGE)
|
|
tBsB_s2r = thr_s2r_b.partition_S(smem_bt)
|
|
|
|
# ((S2R_ATOM_V, S2R_REST_V), S2R_M, S2R_N)
|
|
tBrB_s2r = cute.make_rmem_tensor(
|
|
cute.slice_(tBsB_s2r.shape, (None, None, None, 0)),
|
|
dtype,
|
|
)
|
|
return tiled_s2r_b, tBsB_s2r, tBrB_s2r
|
|
|
|
def pre_inter_smem_store_and_partition_b(
|
|
self, local_tidx, smem_bt_internal, tiled_s2r_b, tBrB_s2r
|
|
):
|
|
dtype = smem_bt_internal.element_type
|
|
# Make tiledCopy from register to smem store Bt
|
|
copy_atom_r2s_b = cute.make_copy_atom(
|
|
cute.nvgpu.CopyUniversalOp(),
|
|
dtype,
|
|
num_bits_per_copy=128,
|
|
)
|
|
tiled_r2s_b = cute.make_tiled_copy_S(copy_atom_r2s_b, tiled_s2r_b)
|
|
thr_r2s_b = tiled_r2s_b.get_slice(local_tidx)
|
|
|
|
# Partition shared tensor for smem store Bt
|
|
# ((R2S_ATOM_V, R2S_REST_V), R2S_M, R2S_N, INTERNAL_STAGE)
|
|
tBsB_r2s = thr_r2s_b.partition_D(smem_bt_internal)
|
|
|
|
# Make register fragments for smem load/store Bt
|
|
# ((S2R_ATOM_V, S2R_REST_V), S2R_M, S2R_N)
|
|
tBrB_r2s = thr_r2s_b.retile(tBrB_s2r)
|
|
return tiled_r2s_b, tBrB_r2s, tBsB_r2s
|
|
|
|
def smem_load_and_partition_delta_d(
|
|
self, tiled_s2r_b, local_tidx, smem_delta, smem_tile_coord
|
|
):
|
|
dtype = smem_delta.element_type
|
|
s2r_atom_delta = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), dtype)
|
|
|
|
thr_s2r_b = tiled_s2r_b.get_slice(local_tidx)
|
|
# ((S2R_ATOM_V, S2R_REST_V), S2R_M, S2R_N, INPUT_STAGE)
|
|
tBsDelta_s2r = thr_s2r_b.partition_D(smem_delta)
|
|
|
|
# Make register fragments for smem load/store of Delta/DeltaA
|
|
# ((S2R_ATOM_V, S2R_REST_V), S2R_M, S2R_N)
|
|
tBrDelta_s2r = cute.make_rmem_tensor(tBsDelta_s2r[smem_tile_coord].shape, dtype)
|
|
return s2r_atom_delta, tBsDelta_s2r, tBrDelta_s2r
|
|
|
|
def pre_inter_tmem_load_and_partition_p(self, local_tidx, tInter1, smem_pt):
|
|
copy_atom_t2r_inter1 = cute.make_copy_atom(
|
|
tcgen05.Ld16x256bOp(tcgen05.Repetition(8), tcgen05.Pack.NONE),
|
|
self.acc_dtype,
|
|
)
|
|
return self.make_tmem_load_and_partition(
|
|
copy_atom_t2r_inter1,
|
|
tInter1,
|
|
(None, None, 0),
|
|
local_tidx,
|
|
smem_pt[None, None, 0],
|
|
)
|
|
|
|
def make_tmem_load_and_partition(
|
|
self, copy_atom_t2r, tmem_tensor, tmem_tile_coord, local_tidx, smem_tensor
|
|
):
|
|
dtype = tmem_tensor.element_type
|
|
tiled_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tmem_tensor[tmem_tile_coord])
|
|
thr_t2r = tiled_t2r.get_slice(local_tidx)
|
|
# Partition tmem/shared tensor for tmem load INTER1_ACC
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N)
|
|
tTR_t = thr_t2r.partition_S(tmem_tensor)
|
|
tTR_s = thr_t2r.partition_D(smem_tensor)
|
|
# Make register fragments for tmem load INTER1_ACC
|
|
# ((T2R_ATOM_V, T2R_REST_V), T2R_M, T2R_N)
|
|
tTR_r = cute.make_rmem_tensor(
|
|
tTR_s.shape,
|
|
dtype,
|
|
)
|
|
return tiled_t2r, tTR_t, tTR_r
|
|
|
|
def smem_store_and_partition_p_y(self, local_tidx, smem_pt, tiled_t2r_inter1):
|
|
dtype = smem_pt.element_type
|
|
copy_atom_r2s_p = cute.make_copy_atom(
|
|
cute.nvgpu.warp.StMatrix8x8x16bOp(transpose=True, num_matrices=4),
|
|
dtype,
|
|
)
|
|
tiled_r2s_p = cute.make_tiled_copy_D(copy_atom_r2s_p, tiled_t2r_inter1)
|
|
thr_r2s_p = tiled_r2s_p.get_slice(local_tidx)
|
|
# Partition smem/register tensor for smem store INTER2_P
|
|
# ((R2S_ATOM_V, R2S_REST_V), R2S_M, R2S_N, INTERNAL_STAGE)
|
|
tRS_sP = thr_r2s_p.partition_D(smem_pt)
|
|
# ((R2S_ATOM_V, R2S_REST_V), R2S_M, R2S_N)
|
|
tRS_rP = cute.make_rmem_tensor(
|
|
cute.slice_(tRS_sP.shape, (None, None, None, 0)), self.io_dtype
|
|
)
|
|
return tiled_r2s_p, tRS_rP, tRS_sP
|
|
|
|
def pre_inter_make_delta(self, smem_delta, smem_bt_layout):
|
|
# Broadcast Delta/DeltaA to Bt shape on M dimension
|
|
# before: (128,(64,2),2):(64,(1,8192),16384)
|
|
# after : (128,(64,2),2):(0,(1,64),128)
|
|
# (MMA, MMA_M, MMA_K, INPUT_STAGE)
|
|
sDeltaA = cute.make_tensor(
|
|
smem_delta.iterator,
|
|
cute.make_layout(
|
|
smem_bt_layout.shape,
|
|
stride=(
|
|
0,
|
|
(1, cute.get(smem_bt_layout.shape, mode=[1, 0])),
|
|
smem_delta.layout.stride[1],
|
|
),
|
|
),
|
|
)
|
|
return sDeltaA
|
|
|
|
def pre_inter_scale_bt_with_delta(
|
|
self, tBrB_s2r, tBrDelta_s2r, tBrDeltaA_s2r, last_column
|
|
):
|
|
tCompute = cute.make_rmem_tensor(tBrB_s2r.shape, self.acc_dtype)
|
|
tBrB_Compute = cute.make_rmem_tensor(tBrB_s2r.shape, self.acc_dtype)
|
|
tBrDelta_Compute = cute.make_rmem_tensor(tBrDelta_s2r.shape, self.acc_dtype)
|
|
tBrDeltaA_Compute = cute.make_rmem_tensor(tBrDeltaA_s2r.shape, self.acc_dtype)
|
|
|
|
tBrB_Compute.store(tBrB_s2r.load().to(self.acc_dtype))
|
|
tBrDelta_Compute.store(tBrDelta_s2r.load().to(self.acc_dtype))
|
|
tBrDeltaA_Compute.store(tBrDeltaA_s2r.load().to(self.acc_dtype))
|
|
|
|
for reg_idx in range(0, cute.size(tBrB_Compute), 2):
|
|
tCompute[reg_idx], tCompute[reg_idx + 1] = cute.arch.mul_packed_f32x2(
|
|
(
|
|
cute.math.exp(
|
|
(last_column - tBrDeltaA_Compute[reg_idx]), fastmath=True
|
|
),
|
|
cute.math.exp(
|
|
(last_column - tBrDeltaA_Compute[reg_idx + 1]), fastmath=True
|
|
),
|
|
),
|
|
(tBrDelta_Compute[reg_idx], tBrDelta_Compute[reg_idx + 1]),
|
|
)
|
|
tCompute[reg_idx], tCompute[reg_idx + 1] = cute.arch.mul_packed_f32x2(
|
|
(tCompute[reg_idx], tCompute[reg_idx + 1]),
|
|
(tBrB_Compute[reg_idx], tBrB_Compute[reg_idx + 1]),
|
|
)
|
|
return tCompute
|
|
|
|
def epilog_make_delta(self, smem_cumsum_delta):
|
|
# Broadcast cumsum delta from LxINPUT_STAGE to LxDxINPUT_STAGE
|
|
sDeltaA = cute.make_tensor(
|
|
smem_cumsum_delta.iterator,
|
|
cute.make_layout(
|
|
(*self.tile_shape_mnk_inter2[:2], self.input_stages),
|
|
stride=(1, 0, smem_cumsum_delta.layout.shape[0]),
|
|
),
|
|
)
|
|
return sDeltaA
|
|
|
|
def epilog_make_d(self, smem_d):
|
|
# Broadcast d from DxINPUT_STAGE to LxDxINPUT_STAGE
|
|
sD = cute.make_tensor(
|
|
smem_d.iterator,
|
|
cute.make_layout(
|
|
(*self.tile_shape_mnk_inter2[:2], self.input_stages),
|
|
stride=(0, 1, smem_d.layout.shape[0]),
|
|
),
|
|
)
|
|
return sD
|
|
|
|
def epilog_tma_partition_y(self, tma_tensor_y, tma_atom_y, smem_y, epi_tile):
|
|
# Local_tile partition global tensors
|
|
# (L, D, 1, 1, C, EH, B)
|
|
gY = cute.local_tile(
|
|
tma_tensor_y,
|
|
cute.slice_(self.tile_shape_mnk_inter2, (None, None, 0)),
|
|
(None, None, None, None, None),
|
|
)
|
|
# (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, 1, 1, C, EH, B)
|
|
gY_epi = cute.flat_divide(gY, epi_tile)
|
|
# ((ATOM_V, REST_V), INPUT_STAGE)
|
|
# ((ATOM_V, REST_V), EPI_M, EPI_N, 1, 1, C, EH, B)
|
|
bSG_sY, bSG_gY_pre_slice = cpasync.tma_partition(
|
|
tma_atom_y,
|
|
0,
|
|
cute.make_layout(1),
|
|
cute.group_modes(smem_y, 0, 2),
|
|
cute.group_modes(gY_epi, 0, 2),
|
|
)
|
|
return bSG_sY, bSG_gY_pre_slice
|
|
|
|
def epilog_smem_load_and_partition_x(
|
|
self, tiled_t2r_inter2_intra2, local_tidx, smem_xt, epi_tile
|
|
):
|
|
dtype = smem_xt.element_type
|
|
copy_atom_s2r_x = cute.make_copy_atom(
|
|
cute.nvgpu.warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4),
|
|
dtype,
|
|
)
|
|
tiled_s2r_x = cute.make_tiled_copy_D(copy_atom_s2r_x, tiled_t2r_inter2_intra2)
|
|
thr_s2r_x = tiled_s2r_x.get_slice(local_tidx)
|
|
# Partition smem/register tensor for smem store INTER2_P
|
|
# (R2S_ATOM, R2S_M, R2S_N, EPI_M, EPI_N, INPUT_STAGES)
|
|
tSR_sX = thr_s2r_x.partition_S(cute.flat_divide(smem_xt, epi_tile))
|
|
# (R2S_ATOM, R2S_M, R2S_N)
|
|
tSR_rX = cute.make_rmem_tensor(
|
|
cute.slice_(tSR_sX.shape, (None, None, None, 0, 0, 0)), dtype
|
|
)
|
|
return tiled_s2r_x, tSR_sX, tSR_rX
|
|
|
|
def epilog_tmem_load_and_partition_acc(self, local_tidx, tIntra, smem_y):
|
|
copy_atom_t2r_inter2_intra2 = cute.make_copy_atom(
|
|
tcgen05.Ld16x256bOp(tcgen05.Repetition(4), tcgen05.Pack.NONE),
|
|
self.acc_dtype,
|
|
)
|
|
return self.make_tmem_load_and_partition(
|
|
copy_atom_t2r_inter2_intra2,
|
|
tIntra,
|
|
(None, None, 0, 0, 0),
|
|
local_tidx,
|
|
smem_y[None, None, 0],
|
|
)
|
|
|
|
|
|
def run(
|
|
gbehcdln: Tuple[int, int, int, int, int, int, int, int],
|
|
io_dtype: Type[cutlass.Numeric],
|
|
cumsum_delta_dtype: Type[cutlass.Numeric],
|
|
acc_dtype: Type[cutlass.Numeric],
|
|
fuse_scale_d: str,
|
|
tolerance: float,
|
|
print_rtol_stats: bool,
|
|
ref_lower_precision: bool,
|
|
warmup_iterations: int,
|
|
iterations: int,
|
|
skip_ref_check: bool,
|
|
use_cold_l2: bool = False,
|
|
**kwargs,
|
|
):
|
|
has_d = fuse_scale_d != "none"
|
|
d_has_hdim = fuse_scale_d == "vector"
|
|
|
|
print("Running B100 Mamba2 SSD with:")
|
|
print(f"GBEHCDLN: {gbehcdln}")
|
|
print(
|
|
f"Input/Output dtype: {io_dtype}, Intermediate delta dtype: {cumsum_delta_dtype}, Acc dtype: {acc_dtype}"
|
|
)
|
|
print(
|
|
f"Has D (True means fuse Y+=X*D): {has_d}, D has Hdim (True means D.shape DxEH, False means 1xEH): {d_has_hdim}"
|
|
)
|
|
print(f"Tolerance: {tolerance}")
|
|
print(f"Warmup iterations: {warmup_iterations}")
|
|
print(f"Iterations: {iterations}")
|
|
print(f"Skip reference checking: {skip_ref_check}")
|
|
print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}")
|
|
|
|
# Unpack parameters
|
|
G, B, E, H, C, D, L, N = gbehcdln
|
|
EH = E * H
|
|
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("GPU is required to run this example!")
|
|
|
|
# Match same seed in ssd_reference.py for reference check
|
|
torch.manual_seed(42)
|
|
|
|
# Create and permute tensor A/B/C
|
|
def create_and_permute_tensor(
|
|
shape, permute_order, dtype, dt_or_a=0, dynamic_modes=None, ref_tensor=None
|
|
):
|
|
# Build fp32 reference torch tensor
|
|
if ref_tensor is None:
|
|
ref_tensor = (
|
|
torch.empty(*shape, dtype=torch.float32)
|
|
# .random_(-1, 1)
|
|
.normal_(0, 0.5)
|
|
# .uniform_(-1,1)
|
|
.permute(permute_order)
|
|
)
|
|
if dt_or_a == 1: # dt:
|
|
ref_tensor = F.softplus(ref_tensor - 4)
|
|
elif dt_or_a == 2: # A:
|
|
ref_tensor = -torch.exp(ref_tensor)
|
|
|
|
# Build torch_dtype torch tensor
|
|
torch_dtype = cutlass_torch.dtype(dtype)
|
|
|
|
dst_tensor = ref_tensor.to(dtype=torch_dtype).cuda()
|
|
cute_tensor = from_dlpack(dst_tensor, assumed_align=16)
|
|
for mode in dynamic_modes:
|
|
cute_tensor = cute_tensor.mark_compact_shape_dynamic(
|
|
mode=mode, stride_order=dst_tensor.dim_order()
|
|
)
|
|
|
|
return ref_tensor, cute_tensor, dst_tensor
|
|
|
|
# INPUT tensors
|
|
# x: (D, L, C, EH, B):(C*L, 1, L, D*C*L, EH*D*C*L)
|
|
x_ref, x_tensor, x_torch = create_and_permute_tensor(
|
|
[B, EH, D, C, L], [2, 4, 3, 1, 0], io_dtype, dynamic_modes=[2, 3, 4]
|
|
)
|
|
# delta/delta_a/cumsum_delta: (L, C, EH, B):(1, L, C*L, EH*C*L)
|
|
delta_ref, delta_tensor, delta_torch = create_and_permute_tensor(
|
|
[B, EH, C, L], [3, 2, 1, 0], io_dtype, dt_or_a=1, dynamic_modes=[1, 2, 3]
|
|
)
|
|
# a: (EH):(1)
|
|
a_ref, a_tensor, a_torch = create_and_permute_tensor(
|
|
[EH], [0], io_dtype, dt_or_a=2, dynamic_modes=[0]
|
|
)
|
|
|
|
if has_d:
|
|
# d: (D, EH):(1, D) or (1, EH):(0, 1)
|
|
d_ref, d_tensor, d_torch = create_and_permute_tensor(
|
|
[EH, D if d_has_hdim else 1], [1, 0], io_dtype, dynamic_modes=[1]
|
|
)
|
|
else:
|
|
d_ref = None
|
|
d_tensor = None
|
|
|
|
# b/c: (L, N, C, G, B):(1, C*L, L, N*C*L, G*N*C*L)
|
|
b_ref, b_tensor, b_torch = create_and_permute_tensor(
|
|
[B, G, N, C, L], [4, 2, 3, 1, 0], io_dtype, dynamic_modes=[2, 3, 4]
|
|
)
|
|
c_ref, c_tensor, c_torch = create_and_permute_tensor(
|
|
[B, G, N, C, L], [4, 2, 3, 1, 0], io_dtype, dynamic_modes=[2, 3, 4]
|
|
)
|
|
|
|
# OUTPUT tensors
|
|
# y: (L, D, C, EH, B):(1, C*L, L, D*C*L, EH*D*C*L)
|
|
y_ref, y_tensor, y_torch = create_and_permute_tensor(
|
|
[B, EH, D, C, L], [4, 2, 3, 1, 0], io_dtype, dynamic_modes=[2, 3, 4]
|
|
)
|
|
# fstate: (D, N, EH, B):(N, 1, D*N, EH*D*N)
|
|
fstate_ref, fstate_tensor, fstate_torch = create_and_permute_tensor(
|
|
[B, EH, D, N], [2, 3, 1, 0], io_dtype, dynamic_modes=[2, 3]
|
|
)
|
|
|
|
# Call pytorch reference on cpu
|
|
if not ref_lower_precision:
|
|
ssd_reference_fp32_all(
|
|
x_ref,
|
|
a_ref,
|
|
delta_ref,
|
|
b_ref,
|
|
c_ref,
|
|
y_ref,
|
|
fstate_ref,
|
|
d_ref,
|
|
has_d,
|
|
d_has_hdim,
|
|
)
|
|
else:
|
|
ssd_reference_lowprecision_intermediates(
|
|
x_ref,
|
|
a_ref,
|
|
delta_ref,
|
|
b_ref,
|
|
c_ref,
|
|
y_ref,
|
|
fstate_ref,
|
|
cutlass_torch.dtype(io_dtype),
|
|
d_ref,
|
|
has_d,
|
|
d_has_hdim,
|
|
)
|
|
|
|
# Compute cumsum with pytorch on cpu
|
|
delta_a_ref = delta_ref * a_ref.view(1, 1, -1, 1)
|
|
cumsum_delta_ref = torch.empty([B, EH, C, L], dtype=torch.float32).permute(
|
|
[3, 2, 1, 0]
|
|
)
|
|
cumsum_delta_ref.copy_(torch.cumsum(delta_a_ref, dim=0).permute([0, 1, 2, 3]))
|
|
# Copy cumsum_delta_ref to cumsum_delta_tensor
|
|
(
|
|
cumsum_delta_ref,
|
|
cumsum_delta_tensor,
|
|
cumsum_delta_torch,
|
|
) = create_and_permute_tensor(
|
|
[B, EH, C, L],
|
|
[3, 2, 1, 0],
|
|
cumsum_delta_dtype,
|
|
ref_tensor=cumsum_delta_ref,
|
|
dynamic_modes=[1, 2, 3],
|
|
)
|
|
|
|
# Call fused ssd kernel
|
|
ssd = SSDKernel(
|
|
io_dtype,
|
|
cumsum_delta_dtype,
|
|
acc_dtype,
|
|
L,
|
|
D,
|
|
N,
|
|
has_d,
|
|
d_has_hdim,
|
|
)
|
|
|
|
# Compute max active clusters on current device
|
|
hardware_info = cutlass.utils.HardwareInfo()
|
|
max_active_clusters = hardware_info.get_max_active_clusters(1)
|
|
|
|
stream = cutlass.cuda.default_stream()
|
|
# Compile ssd kernel
|
|
compiled_ssd = cute.compile(
|
|
ssd,
|
|
x_tensor,
|
|
cumsum_delta_tensor,
|
|
delta_tensor,
|
|
b_tensor,
|
|
c_tensor,
|
|
y_tensor,
|
|
fstate_tensor,
|
|
d_tensor,
|
|
max_active_clusters,
|
|
stream,
|
|
)
|
|
|
|
# Launch compiled ssd kernel for reference check
|
|
if not skip_ref_check:
|
|
compiled_ssd(
|
|
x_tensor,
|
|
cumsum_delta_tensor,
|
|
delta_tensor,
|
|
b_tensor,
|
|
c_tensor,
|
|
y_tensor,
|
|
fstate_tensor,
|
|
d_tensor,
|
|
stream,
|
|
)
|
|
|
|
# Reference check
|
|
if print_rtol_stats:
|
|
print("\nY's Relative diffs:")
|
|
analyze_relative_diffs(
|
|
y_torch.cpu(), y_ref.to(cutlass_torch.dtype(io_dtype))
|
|
)
|
|
print("\nFstate's Relative diffs:")
|
|
analyze_relative_diffs(
|
|
fstate_torch.cpu(), fstate_ref.to(cutlass_torch.dtype(io_dtype))
|
|
)
|
|
torch.testing.assert_close(
|
|
y_torch.cpu(),
|
|
y_ref.to(cutlass_torch.dtype(io_dtype)),
|
|
atol=tolerance,
|
|
rtol=1e-02,
|
|
)
|
|
torch.testing.assert_close(
|
|
fstate_torch.cpu(),
|
|
fstate_ref.to(cutlass_torch.dtype(io_dtype)),
|
|
atol=tolerance,
|
|
rtol=1e-05,
|
|
)
|
|
|
|
def generate_tensors():
|
|
# Reuse existing CPU reference tensors and create new GPU tensors from them
|
|
_, x_tensor_new, _ = create_and_permute_tensor(
|
|
[B, EH, D, C, L],
|
|
[2, 4, 3, 1, 0],
|
|
io_dtype,
|
|
ref_tensor=x_ref,
|
|
dynamic_modes=[2, 3, 4],
|
|
)
|
|
_, cumsum_delta_tensor_new, _ = create_and_permute_tensor(
|
|
[B, EH, C, L],
|
|
[3, 2, 1, 0],
|
|
cumsum_delta_dtype,
|
|
ref_tensor=cumsum_delta_ref,
|
|
dynamic_modes=[1, 2, 3],
|
|
)
|
|
_, delta_tensor_new, _ = create_and_permute_tensor(
|
|
[B, EH, C, L],
|
|
[3, 2, 1, 0],
|
|
io_dtype,
|
|
ref_tensor=delta_ref,
|
|
dynamic_modes=[1, 2, 3],
|
|
)
|
|
_, b_tensor_new, _ = create_and_permute_tensor(
|
|
[B, G, N, C, L],
|
|
[4, 2, 3, 1, 0],
|
|
io_dtype,
|
|
ref_tensor=b_ref,
|
|
dynamic_modes=[2, 3, 4],
|
|
)
|
|
_, c_tensor_new, _ = create_and_permute_tensor(
|
|
[B, G, N, C, L],
|
|
[4, 2, 3, 1, 0],
|
|
io_dtype,
|
|
ref_tensor=c_ref,
|
|
dynamic_modes=[2, 3, 4],
|
|
)
|
|
_, y_tensor_new, _ = create_and_permute_tensor(
|
|
[B, EH, D, C, L],
|
|
[4, 2, 3, 1, 0],
|
|
io_dtype,
|
|
ref_tensor=y_ref,
|
|
dynamic_modes=[2, 3, 4],
|
|
)
|
|
_, fstate_tensor_new, _ = create_and_permute_tensor(
|
|
[B, EH, D, N],
|
|
[2, 3, 1, 0],
|
|
io_dtype,
|
|
ref_tensor=fstate_ref,
|
|
dynamic_modes=[2, 3],
|
|
)
|
|
|
|
if has_d:
|
|
_, d_tensor_new, _ = create_and_permute_tensor(
|
|
[EH, D if d_has_hdim else 1],
|
|
[1, 0],
|
|
io_dtype,
|
|
ref_tensor=d_ref,
|
|
dynamic_modes=[1],
|
|
)
|
|
else:
|
|
d_tensor_new = d_tensor
|
|
|
|
return testing.JitArguments(
|
|
x_tensor_new,
|
|
cumsum_delta_tensor_new,
|
|
delta_tensor_new,
|
|
b_tensor_new,
|
|
c_tensor_new,
|
|
y_tensor_new,
|
|
fstate_tensor_new,
|
|
d_tensor_new,
|
|
stream,
|
|
)
|
|
|
|
workspace_count = 1
|
|
if use_cold_l2:
|
|
one_workspace_bytes = (
|
|
x_torch.numel() * x_torch.element_size()
|
|
+ cumsum_delta_torch.numel() * cumsum_delta_torch.element_size()
|
|
+ delta_torch.numel() * delta_torch.element_size()
|
|
+ b_torch.numel() * b_torch.element_size()
|
|
+ c_torch.numel() * c_torch.element_size()
|
|
+ y_torch.numel() * y_torch.element_size()
|
|
+ fstate_torch.numel() * fstate_torch.element_size()
|
|
)
|
|
if has_d:
|
|
one_workspace_bytes += d_torch.numel() * d_torch.element_size()
|
|
|
|
workspace_count = testing.get_workspace_count(
|
|
one_workspace_bytes, warmup_iterations, iterations
|
|
)
|
|
|
|
exec_time = testing.benchmark(
|
|
compiled_ssd,
|
|
workspace_generator=generate_tensors,
|
|
workspace_count=workspace_count,
|
|
stream=stream,
|
|
warmup_iterations=warmup_iterations,
|
|
iterations=iterations,
|
|
)
|
|
|
|
return exec_time # Return execution time in microseconds
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
def parse_comma_separated_ints(s: str) -> List[int]:
|
|
try:
|
|
return [int(x.strip()) for x in s.split(",")]
|
|
except ValueError:
|
|
raise argparse.ArgumentTypeError(
|
|
"Invalid format. Expected comma-separated integers."
|
|
)
|
|
|
|
parser = argparse.ArgumentParser(
|
|
description="Example of MxNxKxL GEMM on Blackwell."
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--gbehcdln",
|
|
type=parse_comma_separated_ints,
|
|
default=[2, 4, 2, 40, 32, 64, 128, 128],
|
|
# default=[2, 3, 2, 2, 8, 64, 128, 128],
|
|
# default=[1, 2, 1, 4, 8, 64, 128, 128],
|
|
help="gbehcdln dimensions (comma-separated)",
|
|
)
|
|
parser.add_argument("--io_dtype", type=cutlass.dtype, default=cutlass.BFloat16)
|
|
parser.add_argument(
|
|
"--cumsum_delta_dtype", type=cutlass.dtype, default=cutlass.Float32
|
|
)
|
|
parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32)
|
|
parser.add_argument(
|
|
"--fuse_scale_d",
|
|
type=str,
|
|
choices=["none", "scalar", "vector"],
|
|
default="vector",
|
|
help="Fuse scale type: none (no Y+=X*D fusion), scalar (Y+=X*D fusion with D.shape=1xEH), or vector (Y+=X*D fusion with D.shape=DxEH)",
|
|
)
|
|
parser.add_argument(
|
|
"--ref_lower_precision",
|
|
action="store_true",
|
|
default=True,
|
|
help="Use lower precision for reference check",
|
|
)
|
|
parser.add_argument(
|
|
"--no-ref_lower_precision",
|
|
action="store_false",
|
|
dest="ref_lower_precision",
|
|
default=False,
|
|
help="Disable lower precision for reference check",
|
|
)
|
|
parser.add_argument(
|
|
"--tolerance", type=float, default=5e-02, help="Tolerance for validation"
|
|
)
|
|
parser.add_argument(
|
|
"--print_rtol_stats",
|
|
action="store_true",
|
|
default=True,
|
|
help="Enable print rtol stats",
|
|
)
|
|
parser.add_argument(
|
|
"--no-print_rtol_stats",
|
|
action="store_false",
|
|
dest="print_rtol_stats",
|
|
default=False,
|
|
help="Disable print rtol stats",
|
|
)
|
|
parser.add_argument(
|
|
"--warmup_iterations",
|
|
type=int,
|
|
default=0,
|
|
help="Number of warmup iterations",
|
|
)
|
|
parser.add_argument(
|
|
"--iterations",
|
|
type=int,
|
|
default=1,
|
|
help="Number of iterations",
|
|
)
|
|
parser.add_argument(
|
|
"--skip_ref_check", action="store_true", help="Skip reference checking"
|
|
)
|
|
parser.add_argument(
|
|
"--use_cold_l2",
|
|
action="store_true",
|
|
default=False,
|
|
help="Use circular buffer tensor sets to ensure L2 cold cache",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
if len(args.gbehcdln) != 8:
|
|
parser.error("--gbehcdln must contain exactly 8 values")
|
|
|
|
run(
|
|
args.gbehcdln,
|
|
args.io_dtype,
|
|
args.cumsum_delta_dtype,
|
|
args.acc_dtype,
|
|
args.fuse_scale_d,
|
|
args.tolerance,
|
|
args.print_rtol_stats,
|
|
args.ref_lower_precision,
|
|
args.warmup_iterations,
|
|
args.iterations,
|
|
args.skip_ref_check,
|
|
args.use_cold_l2,
|
|
)
|
|
print("PASS")
|