ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-11 17:51:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

moe-gemm change into flatmm & tile_window modify to tile_scatter_gather

Browse Source
This commit is contained in:
zanzhang
2025-04-27 15:55:27 +08:00
parent b981730d9a
commit af2e0e51fc
23 changed files with 1957 additions and 844 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
example/ck_tile/XX_moe_gemm/moe_gemm.hpp
View File

@@ -7,7 +7,7 @@
#include "ck_tile/core.hpp"
#include "ck_tile/host/kernel_launch.hpp"
#include "ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp"
#include "ck_tile/ops/moe_gemm.hpp"
template <typename DataType>
struct GemmTypeConfig;
@@ -37,6 +37,9 @@ auto create_args(int argc, char* argv[])
arg_parser.insert("experts", "8", "Num of experts - 8 by default")
.insert("NumTokens", "128", "M dimensions - 128 by default.")
.insert("TopK", "3", "Top K - 2 by default.")
// .insert("TopK", "2", "Top K - 2 by default.")
// .insert("N", "8192", "N dimensions - 4096 by default.")
// .insert("K", "6144", "K dimensions - 4096 by default.")
.insert("N", "4096", "N dimensions - 4096 by default.")
.insert("K", "4096", "K dimensions - 4096 by default.")
.insert("stride_A", "", "Tensor A strides - it is empty by default.")
@@ -46,6 +49,7 @@ auto create_args(int argc, char* argv[])
.insert("b_layout", "C", "B tensor data layout - Col by default.")
.insert("c_layout", "R", "C tensor data layout - Row by default.")
.insert("validate", "1", "0. No validation, 1. Validation on CPU.")
.insert("prec", "fp16", "data type. fp16/bf16/fp8/bf8")
.insert("repeat", "10", "number of iterations to benchmark the kernel.");
bool result = arg_parser.parse(argc, argv);

138
example/ck_tile/XX_moe_gemm/moe_gemm1_xdl_fp8.cpp
View File

@@ -13,12 +13,11 @@
#include "ck_tile/core.hpp"
#include "ck_tile/ops/epilogue.hpp"
#include "ck_tile/ops/gemm.hpp"
#include "ck_tile/ops/flatmm.hpp"
#include "ck_tile/host.hpp"
#include "moe_gemm.hpp"
#include "ck_tile/host/reference/reference_fused_single_moe_gemm.hpp"
namespace {
struct MoeGemmKernelParam
{
static const bool kPadM = false;
@@ -30,8 +29,8 @@ struct MoeGemmKernelParam
static const ck_tile::index_t N_Tile = 128;
static const ck_tile::index_t K_Tile = 32; // need to ensure the M_per_thread = 1
static const ck_tile::index_t M_Warp = 2;
static const ck_tile::index_t N_Warp = 2;
static const ck_tile::index_t M_Warp = 1;
static const ck_tile::index_t N_Warp = 4;
static const ck_tile::index_t K_Warp = 1;
static const ck_tile::index_t M_Warp_Tile = 32;
@@ -39,85 +38,64 @@ struct MoeGemmKernelParam
static const ck_tile::index_t K_Warp_Tile = 16;
};
using CodegenGemmShape = ck_tile::TileGemmShape<ck_tile::sequence<MoeGemmKernelParam::M_Tile,
MoeGemmKernelParam::N_Tile,
MoeGemmKernelParam::K_Tile>,
ck_tile::sequence<MoeGemmKernelParam::M_Warp,
MoeGemmKernelParam::N_Warp,
MoeGemmKernelParam::K_Warp>,
ck_tile::sequence<MoeGemmKernelParam::M_Warp_Tile,
MoeGemmKernelParam::N_Warp_Tile,
MoeGemmKernelParam::K_Warp_Tile>>;
using TilePartitioner = ck_tile::GemmTile1DPartitioner<CodegenGemmShape>;
template <typename ALayout, typename BLayout, typename CLayout>
using CodegenGemmTraits = ck_tile::TileGemmUniversalTraits<MoeGemmKernelParam::kPadM,
MoeGemmKernelParam::kPadN,
MoeGemmKernelParam::kPadK,
true,
ALayout,
BLayout,
CLayout>;
template <typename ALayout, typename BLayout, typename CLayout>
using CodegenPipelineProblem =
ck_tile::UniversalGemmPipelineProblem<ADataType,
BDataType,
AccDataType,
CodegenGemmShape,
CodegenGemmTraits<ALayout, BLayout, CLayout>>;
template <typename ALayout, typename BLayout, typename CLayout>
using CodegenGemmPipeline =
ck_tile::MoeGemmPipelineAgBgCrImpl<CodegenPipelineProblem<ALayout, BLayout, CLayout>>;
template <typename ALayout, typename BLayout, typename CLayout>
using GemmEpilogue = ck_tile::CShuffleEpilogue<ck_tile::CShuffleEpilogueProblem<
ADataType,
BDataType,
AccDataType,
CDataType,
CLayout,
CodegenPipelineProblem<ALayout, BLayout, CLayout>::kBlockSize,
TilePartitioner::MPerBlock,
TilePartitioner::NPerBlock,
MoeGemmKernelParam::M_Warp,
MoeGemmKernelParam::N_Warp,
MoeGemmKernelParam::M_Warp_Tile,
MoeGemmKernelParam::N_Warp_Tile,
MoeGemmKernelParam::K_Warp_Tile,
CodegenPipelineProblem<ALayout, BLayout, CLayout>::TransposeC>>;
// template <typename ALayout, typename BLayout, typename CLayout>
// using GemmEpilogue = ck_tile::DefaultGemm2DEpilogue<ck_tile::DefaultGemm2DEpilogueProblem<
// AccDataType,
// CDataType,
// CLayout,
// MoeGemmKernelParam::kPadM,
// MoeGemmKernelParam::kPadN,
// MoeGemmKernelParam::M_Warp_Tile,
// MoeGemmKernelParam::N_Warp_Tile,
// MoeGemmKernelParam::K_Warp_Tile,
// CodegenPipelineProblem<ALayout, BLayout, CLayout>::TransposeC
// >>;
template <typename ALayout, typename BLayout, typename CLayout>
using Kernel = ck_tile::MoeGemmKernel<TilePartitioner,
CodegenGemmPipeline<ALayout, BLayout, CLayout>,
GemmEpilogue<ALayout, BLayout, CLayout>>;
}; // namespace
template <typename ALayout, typename BLayout, typename CLayout>
float moe_gemm(const moe_gemm_kargs& gemm_desc, const ck_tile::stream_config& s)
{
using MoeGemmKernel = ::Kernel<ALayout, BLayout, CLayout>;
using CodegenMoeGemmShape = ck_tile::TileFlatmmShape<
ck_tile::sequence<MoeGemmKernelParam::M_Tile,
MoeGemmKernelParam::N_Tile,
MoeGemmKernelParam::K_Tile>,
ck_tile::sequence<MoeGemmKernelParam::M_Warp,
MoeGemmKernelParam::N_Warp,
MoeGemmKernelParam::K_Warp>,
ck_tile::sequence<MoeGemmKernelParam::M_Warp_Tile,
MoeGemmKernelParam::N_Warp_Tile,
MoeGemmKernelParam::K_Warp_Tile>>;
using TilePartitioner = ck_tile::GemmTile1DPartitioner<CodegenMoeGemmShape>;
using CodegenMoeGemmTraits = ck_tile::TileGemmTraits<MoeGemmKernelParam::kPadM,
MoeGemmKernelParam::kPadN,
MoeGemmKernelParam::kPadK,
ALayout,
BLayout,
CLayout>;
using CodegenPipelineProblem =
ck_tile::GemmPipelineProblem<ADataType,
BDataType,
AccDataType,
CodegenMoeGemmShape,
CodegenMoeGemmTraits>;
using CodegenMoeGemmPolicy = ck_tile::UniversalFlatmmPipelineAgBgCrPolicy;
using CodegenMoeGemmPipeline =
ck_tile::MoeGemmPipelineAgBgCrImpl<CodegenPipelineProblem, CodegenMoeGemmPolicy>;
using GemmEpilogue = ck_tile::CShuffleEpilogue<ck_tile::CShuffleEpilogueProblem<
ADataType,
BDataType,
AccDataType,
CDataType,
CLayout,
CodegenPipelineProblem::kBlockSize,
TilePartitioner::MPerBlock,
TilePartitioner::NPerBlock,
MoeGemmKernelParam::M_Warp,
MoeGemmKernelParam::N_Warp,
MoeGemmKernelParam::M_Warp_Tile,
MoeGemmKernelParam::N_Warp_Tile,
MoeGemmKernelParam::K_Warp_Tile,
CodegenPipelineProblem::TransposeC>>;
using Kernel = ck_tile::MoeGemmKernel<TilePartitioner,
CodegenMoeGemmPipeline,
GemmEpilogue>;
// TODO: malloc sorted_tokend_ids buffer
const auto arguments = MoeGemmKernel::MoeGemmKernelArgs::MakeKernelArgs(gemm_desc);
const dim3 grids = MoeGemmKernel::GridSize(gemm_desc.M, gemm_desc.N, 1);
constexpr dim3 blocks = MoeGemmKernel::BlockSize();
const auto arguments = Kernel::MakeKernelArgs(gemm_desc);
const dim3 grids = Kernel::GridSize(gemm_desc.M, gemm_desc.N, 1);
constexpr dim3 blocks = Kernel::BlockSize();
// ck_tile::hip_check_error(hipMemcpyWithStream(
// arguments.data(),
@@ -127,7 +105,7 @@ float moe_gemm(const moe_gemm_kargs& gemm_desc, const ck_tile::stream_config& s)
if(s.log_level_ > 0)
{
std::cout << "Launching kernel: " << MoeGemmKernel::GetName() << " with args:"
std::cout << "Launching kernel: " << Kernel::GetName() << " with args:"
<< " grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
<< ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z << "}"
<< std::endl;
@@ -136,7 +114,7 @@ float moe_gemm(const moe_gemm_kargs& gemm_desc, const ck_tile::stream_config& s)
float ave_time =
ck_tile::launch_kernel(s,
ck_tile::make_kernel<blocks.x, MoeGemmKernelParam::kBlockPerCu>(
MoeGemmKernel{}, grids, blocks, 0, arguments));
Kernel{}, grids, blocks, 0, arguments));
return ave_time;
}

128
example/ck_tile/XX_moe_gemm/run_moe_gemm_example.inc
View File

@@ -15,6 +15,40 @@ static constexpr inline auto is_row_major(Layout layout_)
ck_tile::tensor_layout::gemm::RowMajor>>{};
}
template <typename T>
auto shuffle_b(const ck_tile::HostTensor<T>& t, std::string mfma_dtype, int mfma_type = 0)
{
assert(t.get_lengths().size() == 2);
int n_ = t.get_lengths()[1];
int k_ = t.get_lengths()[0];
if((mfma_dtype == "bf16" || mfma_dtype == "fp16") && mfma_type == 0)
{
ck_tile::HostTensor<T> t_view({n_ / 32, 32, k_ / 16, 2, 8});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4});
}
else if((mfma_dtype == "bf16" || mfma_dtype == "fp16") && mfma_type == 1)
{
ck_tile::HostTensor<T> t_view({n_ / 16, 16, k_ / 32, 4, 8});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4});
}
else if((mfma_dtype == "int8" || mfma_dtype == "fp8") && mfma_type == 0)
{
ck_tile::HostTensor<T> t_view({n_ / 32, 32, k_ / 32, 2, 16});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4});
}
else if((mfma_dtype == "int8" || mfma_dtype == "fp8") && mfma_type == 1)
{
ck_tile::HostTensor<T> t_view({n_ / 16, 16, k_ / 64, 4, 16});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4});
}
return t;
}
template <typename ADataType, typename BDataType, typename AccDataType, typename CDataType>
auto calculate_rtol_atol(const ck_tile::index_t K,
const ck_tile::index_t kbatch,
@@ -37,7 +71,7 @@ auto calculate_rtol_atol(const ck_tile::index_t K,
}
template <typename ALayout, typename BLayout, typename CLayout>
float invoke_gemm(int n_warmup, int n_repeat, const moe_gemm_kargs& args)
float invoke_moe_gemm(int n_warmup, int n_repeat, const moe_gemm_kargs& args)
{
float ave_time = moe_gemm<ALayout, BLayout, CLayout>(
args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat});
@@ -98,51 +132,69 @@ int run_moe_gemm_example_with_layouts(int argc,
// TODO: replace the magic declaration
const ck_tile::index_t MPerBlock = 128;
// const ck_tile::index_t max_num_tokens_padded = topk * num_tokens + experts * MPerBlock - topk;
// const ck_tile::index_t max_num_m_blocks = (max_num_tokens_padded + MPerBlock - 1) / MPerBlock;
ck_tile::index_t sorted_tile_num = 8;
ck_tile::index_t valid_tile_num = sorted_tile_num;
// ck_tile::index_t sorted_tile_num = 16;
// ck_tile::index_t valid_tile_num = 13;
ck_tile::index_t sorted_size = sorted_tile_num * MPerBlock;
// ck_tile::index_t valid_size = valid_tile_num * MPerBlock;
const ck_tile::index_t M = sorted_tile_num * MPerBlock;
// const ck_tile::index_t M = max_num_tokens_padded;
std::unique_ptr<ck_tile::DeviceMem> a_m_k_dev_buf;
std::unique_ptr<ck_tile::DeviceMem> b_k_n_dev_buf;
std::unique_ptr<ck_tile::DeviceMem> b_origin_dev_buf;
std::unique_ptr<ck_tile::DeviceMem> b_shuffle_dev_buf;
std::unique_ptr<ck_tile::DeviceMem> c_m_n_dev_buf;
stride_A = ck_tile::get_default_stride(M, N, stride_A, is_row_major(a_layout));
stride_A = ck_tile::get_default_stride(num_tokens, K, stride_A, is_row_major(a_layout));
stride_B = ck_tile::get_default_stride(K, N, stride_B, is_row_major(b_layout));
stride_C = ck_tile::get_default_stride(M, N, stride_C, is_row_major(CLayout{}));
stride_C = ck_tile::get_default_stride(num_tokens * topk, N, stride_C, is_row_major(CLayout{}));
auto a_m_k_tensor = ck_tile::HostTensor<ADataType>(
ck_tile::host_tensor_descriptor(M, K, stride_A, is_row_major(a_layout)));
ck_tile::host_tensor_descriptor(num_tokens, K, stride_A, is_row_major(a_layout)));
// TODO: add the experts' weights in b
auto b_k_n_tensor = ck_tile::HostTensor<BDataType>(
is_row_major(b_layout)
? ck_tile::host_tensor_descriptor(experts * K, N, stride_B, is_row_major(b_layout))
? ck_tile::host_tensor_descriptor(experts * N, K, stride_B, is_row_major(b_layout))
: ck_tile::host_tensor_descriptor(K, experts * N, stride_B, is_row_major(b_layout)));
auto c_m_n_tensor = ck_tile::HostTensor<CDataType>(
ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{})));
std::cout << "gemm"
<< " a_m_k: " << a_m_k_tensor.mDesc << " b_k_n: " << b_k_n_tensor.mDesc
<< " c_m_n: " << c_m_n_tensor.mDesc << std::endl;
std::string mfma = arg_parser.get_str("prec");
auto c_m_n_tensor = ck_tile::HostTensor<CDataType>(
ck_tile::host_tensor_descriptor(num_tokens * topk, N, stride_C, is_row_major(CLayout{})));
ck_tile::FillUniformDistribution<ADataType>{-1.f, 1.f}(a_m_k_tensor);
ck_tile::FillUniformDistribution<BDataType>{-1.f, 1.f}(b_k_n_tensor);
auto b_shuffle_host = shuffle_b(b_k_n_tensor, mfma, 0);
std::cout << "gemm"
<< " a_m_k: " << a_m_k_tensor.mDesc << " b_k_n: " << b_k_n_tensor.mDesc
<< " b_shuffle: " << b_shuffle_host.mDesc << " c_m_n: " << c_m_n_tensor.mDesc << std::endl;
a_m_k_dev_buf =
std::make_unique<ck_tile::DeviceMem>(a_m_k_tensor.get_element_space_size_in_bytes());
b_k_n_dev_buf =
b_origin_dev_buf =
std::make_unique<ck_tile::DeviceMem>(b_k_n_tensor.get_element_space_size_in_bytes());
b_shuffle_dev_buf =
std::make_unique<ck_tile::DeviceMem>(b_shuffle_host.get_element_space_size_in_bytes());
c_m_n_dev_buf =
std::make_unique<ck_tile::DeviceMem>(c_m_n_tensor.get_element_space_size_in_bytes());
a_m_k_dev_buf->ToDevice(a_m_k_tensor.data());
b_k_n_dev_buf->ToDevice(b_k_n_tensor.data());
b_origin_dev_buf->ToDevice(b_k_n_tensor.data());
b_shuffle_dev_buf->ToDevice(b_shuffle_host.data());
c_m_n_dev_buf->SetZero();
c_m_n_tensor.SetZero();
const void* p_a = a_m_k_dev_buf->GetDeviceBuffer();
const void* p_b = b_k_n_dev_buf->GetDeviceBuffer();
const void* p_b_origin = b_origin_dev_buf->GetDeviceBuffer();
const void* p_b_shuffle = b_shuffle_dev_buf->GetDeviceBuffer();
void* p_c = c_m_n_dev_buf->GetDeviceBuffer();
// TODO: malloc and init sorted tokens and max tokens buffer
@@ -150,7 +202,7 @@ int run_moe_gemm_example_with_layouts(int argc,
ck_tile::HostTensor<ck_tile::index_t> expert_ids(
ck_tile::HostTensorDescriptor({sorted_tile_num}, {1}));
ck_tile::HostTensor<ck_tile::index_t> sorted_token_ids(
ck_tile::HostTensorDescriptor({sorted_tile_num * MPerBlock}, {1}));
ck_tile::HostTensorDescriptor({sorted_size}, {1}));
ck_tile::HostTensor<ck_tile::index_t> max_token_id(
ck_tile::HostTensorDescriptor({1 + sorted_tile_num}));
@@ -163,11 +215,14 @@ int run_moe_gemm_example_with_layouts(int argc,
max_token_id.mData = {valid_tile_num * MPerBlock, 0, 1, 2, 3, 4, 6, 7, 8, 8};
int eids[] = {0, 1, 2, 3, 4, 4, 5, 6, 3, 3, 3, 3}; // {2, 1, 1, 2, 2, 2, 1, 2}
// max_token_id.mData = {valid_size, 2, 2, 1, 1, 2, 2, 2, 2, 2, 2, 1, 2, 2, 0, 0, 0};
// int eids[] = {0, 0, 1, 2, 3, 3, 4, 4, 5, 5, 6, 7, 7, 3, 3, 3};
for(int i = 0; i < sorted_tile_num; i++)
{
expert_ids.mData[i] = eids[i];
}
int token_per_tile = (num_tokens * topk + valid_tile_num - 1) / valid_tile_num;
// int token_per_tile = (num_tokens * topk + valid_tile_num - 1) / valid_tile_num;
int token_per_tile = num_tokens * topk / valid_tile_num;
int tokenid = 0;
// sorted_token_ids.mData[0] = 0;
for(int i = 0; i < sorted_tile_num * MPerBlock; i++)
@@ -199,10 +254,11 @@ int run_moe_gemm_example_with_layouts(int argc,
p_expert_ids_dev,
p_max_token_id_dev,
p_a,
p_b,
p_b_shuffle,
p_c,
num_tokens,
topk,
1, //k_batch
M,
N,
K,
@@ -210,7 +266,7 @@ int run_moe_gemm_example_with_layouts(int argc,
stride_B,
stride_C};
invoke_gemm<ALayout, BLayout, CLayout>(3, repeat, gemm_desc);
invoke_moe_gemm<ALayout, BLayout, CLayout>(3, repeat, gemm_desc);
c_m_n_dev_buf->FromDevice(c_m_n_tensor.data());
@@ -218,7 +274,7 @@ int run_moe_gemm_example_with_layouts(int argc,
if(arg_parser.get_int("validate"))
{
ck_tile::HostTensor<CDataType> c_m_n_host_ref(
ck_tile::host_tensor_descriptor(M, N, stride_C, is_row_major(CLayout{})));
ck_tile::host_tensor_descriptor(num_tokens * topk, N, stride_C, is_row_major(CLayout{})));
c_m_n_host_ref.SetZero();
@@ -238,7 +294,7 @@ int run_moe_gemm_example_with_layouts(int argc,
p_expert_ids_dev,
p_max_token_id_dev,
static_cast<const ADataType*>(p_a),
static_cast<const BDataType*>(p_b),
static_cast<const BDataType*>(p_b_origin),
static_cast<CDataType*>(c_m_n_ref_buf->GetDeviceBuffer()),
num_tokens,
topk,
@@ -254,22 +310,22 @@ int run_moe_gemm_example_with_layouts(int argc,
K, 1 /*kbatch*/, max_accumulated_value);
c_m_n_ref_buf->FromDevice(c_m_n_host_ref.data());
for(int im = 0; im < M; im++)
{
for(int in = 0; in < N; in++)
{
// if (static_cast<float>(static_cast<CDataType*>(p_c)[im * N + in]) != 0)
printf("c[%d][%d]: %f ",
im,
in,
static_cast<float>(static_cast<CDataType*>(p_c)[im * N + in]));
printf("ref[%d][%d]: %f \n",
im,
in,
static_cast<float>(
static_cast<CDataType*>(c_m_n_host_ref.data())[im * N + in]));
}
}
// for(int im = 0; im < M; im++)
// {
// for(int in = 0; in < N; in++)
// {
// // if (static_cast<float>(static_cast<CDataType*>(p_c)[im * N + in]) != 0)
// printf("c[%d][%d]: %f ",
// im,
// in,
// static_cast<float>(static_cast<CDataType*>(p_c)[im * N + in]));
// printf("ref[%d][%d]: %f \n",
// im,
// in,
// static_cast<float>(
// static_cast<CDataType*>(c_m_n_host_ref.data())[im * N + in]));
// }
// }
pass = ck_tile::check_err(c_m_n_tensor,
c_m_n_host_ref,

1
include/ck_tile/core.hpp
View File

@@ -55,6 +55,7 @@
#include "ck_tile/core/tensor/tile_elementwise.hpp"
#include "ck_tile/core/tensor/tile_window.hpp"
#include "ck_tile/core/tensor/tile_window_linear.hpp"
#include "ck_tile/core/tensor/tile_scatter_gather.hpp"
#include "ck_tile/core/tensor/tile_window_utils.hpp"
#include "ck_tile/core/tensor/transpose_tile.hpp"
#include "ck_tile/core/tensor/update_tile.hpp"

2
include/ck_tile/core/tensor/buffer_view.hpp
View File

@@ -5,7 +5,7 @@
#include "ck_tile/core/config.hpp"
#include "ck_tile/core/arch/arch.hpp"
#if __clang_major__ == 20
#if __clang_major__ >= 20
#include "ck_tile/core/arch/amd_buffer_addressing_builtins.hpp"
#else
#include "ck_tile/core/arch/amd_buffer_addressing.hpp"

87
include/ck_tile/core/tensor/load_tile.hpp
View File

@@ -18,32 +18,10 @@
namespace ck_tile {
template <typename BottomTensorView_,
typename WindowLengths_,
typename TileDistribution_,
index_t NumCoord,
template <typename TileWindow_,
index_t i_access = -1,
bool oob_conditional_check = true>
CK_TILE_DEVICE auto load_tile(const tile_window_with_static_distribution<BottomTensorView_,
WindowLengths_,
TileDistribution_,
NumCoord>& tile_window,
number<i_access> = {},
bool_constant<oob_conditional_check> = {})
{
return tile_window.load(number<i_access>{}, bool_constant<oob_conditional_check>{});
}
template <typename BottomTensorView_,
typename WindowLengths_,
typename TileDistribution_,
typename LinearBottomDims_,
index_t i_access = -1,
bool oob_conditional_check = true>
CK_TILE_DEVICE auto load_tile(const tile_window_linear<BottomTensorView_,
WindowLengths_,
TileDistribution_,
LinearBottomDims_>& tile_window,
CK_TILE_DEVICE auto load_tile(const TileWindow_& tile_window,
number<i_access> = {},
bool_constant<oob_conditional_check> = {})
{
@@ -51,35 +29,11 @@ CK_TILE_DEVICE auto load_tile(const tile_window_linear<BottomTensorView_,
}
template <typename DistributedTensor_,
typename BottomTensorView_,
typename WindowLengths_,
typename TileDistribution_,
index_t NumCoord,
typename TileWindow_,
index_t i_access = -1,
bool oob_conditional_check = true>
CK_TILE_DEVICE auto load_tile(DistributedTensor_& dst_tile,
const tile_window_with_static_distribution<BottomTensorView_,
WindowLengths_,
TileDistribution_,
NumCoord>& tile_window,
number<i_access> = {},
bool_constant<oob_conditional_check> = {})
{
return tile_window.load(dst_tile, number<i_access>{}, bool_constant<oob_conditional_check>{});
}
template <typename DistributedTensor_,
typename BottomTensorView_,
typename WindowLengths_,
typename TileDistribution_,
typename LinearBottomDims_,
index_t i_access = -1,
bool oob_conditional_check = true>
CK_TILE_DEVICE auto load_tile(DistributedTensor_& dst_tile,
const tile_window_linear<BottomTensorView_,
WindowLengths_,
TileDistribution_,
LinearBottomDims_>& tile_window,
const TileWindow_& tile_window,
number<i_access> = {},
bool_constant<oob_conditional_check> = {})
{
@@ -138,19 +92,13 @@ CK_TILE_DEVICE auto load_tile_raw(T& tile,
}
template <typename LdsTileWindow_,
typename BottomTensorView_,
typename WindowLengths_,
typename TileDistribution_,
index_t NumCoord,
typename TileWindow_,
index_t i_access = -1,
bool oob_conditional_check = true,
bool pre_nop = false>
CK_TILE_DEVICE auto
async_load_tile_raw(LdsTileWindow_&& lds_tile,
const tile_window_with_static_distribution<BottomTensorView_,
WindowLengths_,
TileDistribution_,
NumCoord>& tile_window,
const TileWindow_& tile_window,
number<i_access> = {},
bool_constant<oob_conditional_check> = {},
bool_constant<pre_nop> = {})
@@ -161,29 +109,6 @@ async_load_tile_raw(LdsTileWindow_&& lds_tile,
bool_constant<pre_nop>{});
}
template <typename LdsTileWindow_,
typename BottomTensorView_,
typename WindowLengths_,
typename TileDistribution_,
typename LinearBottomDims_,
index_t i_access = -1,
bool oob_conditional_check = true,
bool pre_nop = false>
CK_TILE_DEVICE auto async_load_tile_raw(LdsTileWindow_&& lds_tile,
const tile_window_linear<BottomTensorView_,
WindowLengths_,
TileDistribution_,
LinearBottomDims_>& tile_window,
number<i_access> = {},
bool_constant<oob_conditional_check> = {},
bool_constant<pre_nop> = {})
{
return tile_window.async_load_raw(lds_tile,
number<i_access>{},
bool_constant<oob_conditional_check>{},
bool_constant<pre_nop>{});
}
CK_TILE_DEVICE auto async_load_fence(index_t cnt = 0)
{
asm volatile("s_waitcnt vmcnt(%0)" : : "n"(cnt) : "memory");

32
include/ck_tile/core/tensor/tensor_view.hpp
View File

@@ -209,6 +209,26 @@ struct tensor_view
coordinate_has_valid_offset_assuming_top_index_is_valid(desc_, coord),
bool_constant<pre_nop>{});
}
template <typename X,
bool pre_nop = false,
typename std::enable_if<
std::is_same_v<typename vector_traits<remove_cvref_t<X>>::scalar_type,
typename vector_traits<remove_cvref_t<DataType>>::scalar_type>,
bool>::type = false>
CK_TILE_HOST_DEVICE constexpr void
async_get_vectorized_elements_raw(remove_cvref_t<DataType>* smem,
const TensorCoord& coord,
index_t coord_extra_offset,
index_t linear_offset,
bool_constant<pre_nop> = {}) const
{
return buf_.template async_get_raw<X>(
smem,
(coord.get_offset() + coord_extra_offset) / PackedSize,
linear_offset / PackedSize,
coordinate_has_valid_offset_assuming_top_index_is_valid(desc_, coord),
bool_constant<pre_nop>{});
}
template <typename X,
bool pre_nop = false,
@@ -411,21 +431,18 @@ struct null_tensor_view
};
template <address_space_enum BufferAddressSpace = address_space_enum::generic,
amd_buffer_coherence_enum Coherence = amd_buffer_coherence_enum::coherence_default,
typename DataType,
typename... Ts>
CK_TILE_HOST_DEVICE constexpr auto make_tensor_view(DataType* p,
const tensor_descriptor<Ts...>& desc)
{
auto buffer_view =
make_buffer_view<BufferAddressSpace, Coherence>(p, desc.get_element_space_size());
auto buffer_view = make_buffer_view<BufferAddressSpace>(p, desc.get_element_space_size());
return tensor_view<decltype(buffer_view), decltype(desc)>{buffer_view, desc};
}
template <address_space_enum BufferAddressSpace = address_space_enum::generic,
memory_operation_enum DstInMemOp = memory_operation_enum::set,
amd_buffer_coherence_enum Coherence = amd_buffer_coherence_enum::coherence_default,
typename DataType,
typename... Lengths,
typename... Strides,
@@ -444,14 +461,12 @@ make_naive_tensor_view(DataType* p,
number<GuaranteedLastDimensionVectorLength>{},
number<GuaranteedLastDimensionVectorStride>{});
auto buffer_view =
make_buffer_view<BufferAddressSpace, Coherence>(p, desc.get_element_space_size());
auto buffer_view = make_buffer_view<BufferAddressSpace>(p, desc.get_element_space_size());
return tensor_view<decltype(buffer_view), decltype(desc), DstInMemOp>{buffer_view, desc};
}
template <address_space_enum BufferAddressSpace = address_space_enum::generic,
amd_buffer_coherence_enum Coherence = amd_buffer_coherence_enum::coherence_default,
typename DataType,
typename... Lengths,
index_t GuaranteedLastDimensionVectorLength = -1>
@@ -463,8 +478,7 @@ make_naive_tensor_view_packed(DataType* p,
auto desc =
make_naive_tensor_descriptor_packed(lengths, number<GuaranteedLastDimensionVectorLength>{});
auto buffer_view =
make_buffer_view<BufferAddressSpace, Coherence>(p, desc.get_element_space_size());
auto buffer_view = make_buffer_view<BufferAddressSpace>(p, desc.get_element_space_size());
return tensor_view<decltype(buffer_view), decltype(desc)>{buffer_view, desc};
}

735
include/ck_tile/core/tensor/tile_scatter_gather.hpp Normal file
View File

@@ -0,0 +1,735 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck_tile/core/arch/utility.hpp"
#include "ck_tile/core/algorithm/space_filling_curve.hpp"
#include "ck_tile/core/config.hpp"
#include "ck_tile/core/container/array.hpp"
#include "ck_tile/core/container/sequence.hpp"
#include "ck_tile/core/container/tuple.hpp"
#include "ck_tile/core/container/container_helper.hpp"
#include "ck_tile/core/tensor/static_distributed_tensor.hpp"
#include "ck_tile/core/tensor/tensor_adaptor.hpp"
#include "ck_tile/core/tensor/tile_distribution.hpp"
#include "ck_tile/core/utility/functional.hpp"
#include "ck_tile/core/utility/type_traits.hpp"
namespace ck_tile {
/**
* @brief This class provides tile (windowed) view and access to the device memory.
*
* @note This tile window does not support single issue you need to use tile_window_linear
* structure for this purpose
*
* @tparam BottomTensorView_ Class describing & holding device tensor memory.
* @tparam WindowLengths_ Spatial sizes of windowed view on tensor.
* @tparam StaticTileDistribution_ Thread distribution (mapping) into Tile dimensions
* @tparam NumCoord TBD
*/
template <typename BottomTensorView_,
typename WindowLengths_,
typename StaticTileDistribution_,
typename StaticPageIndexArray_,
index_t HsGatherDim = 0,
index_t NumCoord = 1,
index_t YsGatherDim = 0>
struct tile_scatter_gather
{
using BottomTensorView = remove_reference_t<BottomTensorView_>;
using WindowLengths = remove_cvref_t<WindowLengths_>;
using TileDstr = remove_cvref_t<StaticTileDistribution_>;
using PageIdxArray = remove_cvref_t<StaticPageIndexArray_>;
using WindowAdaptor = typename TileDstr::PsYs2XsAdaptor;
using BottomTensorDesc = typename BottomTensorView::TensorDesc;
using DataType = remove_cvref_t<typename BottomTensorView::DataType>;
static constexpr index_t NDimWindowAdaptorTop = WindowAdaptor::get_num_of_top_dimension();
static constexpr index_t NDimBottomTensor = BottomTensorDesc::get_num_of_dimension();
static constexpr index_t NDimP = TileDstr::get_num_of_dimension_p();
static constexpr index_t NDimY = TileDstr::get_num_of_dimension_y();
static constexpr auto I0 = number<0>{};
static constexpr auto I1 = number<1>{};
static_assert(NumCoord == 1);
// TODO: check WindowLengths and StaticTileDistribution are consistent
static_assert(ck_tile::is_known_at_compile_time<WindowLengths>::value,
"wrong! lengths should be static");
static_assert(TileDstr::is_static(), "wrong!");
static_assert(NDimBottomTensor == WindowAdaptor::get_num_of_bottom_dimension(),
"wrong! inconsistent # of diemsnions");
using AdaptorTopIndex = array<index_t, NDimWindowAdaptorTop>;
using BottomTensorIndex = array<index_t, NDimBottomTensor>;
using WindowAdaptorCoord =
decltype(make_tensor_adaptor_coordinate(WindowAdaptor{}, AdaptorTopIndex{}));
using BottomTensorCoord =
decltype(make_tensor_coordinate(BottomTensorDesc{}, BottomTensorIndex{}));
struct load_store_traits
{
private:
static constexpr auto get_vector_dim_y_scalar_per_vector()
{
const auto [ys_vector_lengths, ys_vector_strides] =
tile_scatter_gather::
get_window_adaptor_ys_safe_vector_length_strides();
index_t VectorDimY_ = 0;
index_t ScalarPerVector_ = 1;
for(index_t i = 0; i < NDimY; ++i)
{
if(ys_vector_strides[i] == 1 && ys_vector_lengths[i] > ScalarPerVector_)
{
ScalarPerVector_ = ys_vector_lengths[i];
VectorDimY_ = i;
}
}
return make_tuple(VectorDimY_, ScalarPerVector_);
}
public:
static constexpr index_t PackedSize =
ck_tile::numeric_traits<remove_cvref_t<DataType>>::PackedSize;
static constexpr index_t VectorDimY = get_vector_dim_y_scalar_per_vector().template at<0>();
static constexpr index_t ScalarPerVector =
get_vector_dim_y_scalar_per_vector().template at<1>();
// using vector_type_t = vector_type_maker_t<DataType, ScalarPerVector>;
// using vector_t = typename vector_type_t::type;
using vector_t = thread_buffer<DataType, ScalarPerVector / PackedSize>;
private:
static constexpr auto scalars_per_access_ = [] {
constexpr auto scalars_per_access_arr = generate_array(
[&](auto i) { return (i == VectorDimY) ? ScalarPerVector : 1; }, number<NDimY>{});
/// TODO: add non-automatic storage argument support to macro TO_SEQUENCE()
constexpr auto NDimY_ = NDimY;
return TO_SEQUENCE(scalars_per_access_arr, NDimY_);
}();
static constexpr auto get_space_filling_curve()
{
constexpr auto tile_dstr = TileDstr{};
constexpr auto thread_tensor_lengths_ys =
to_sequence(tile_dstr.get_ys_to_d_descriptor().get_lengths());
// FIXME: need logic to judge dim access order
using DimAccessOrder = typename arithmetic_sequence_gen<0, NDimY, 1>::type;
return space_filling_curve<decltype(thread_tensor_lengths_ys),
DimAccessOrder,
decltype(scalars_per_access_)>{};
}
public:
using SFC_Ys = decltype(get_space_filling_curve());
static constexpr index_t NumAccess = SFC_Ys::get_num_of_access();
static_assert(0 < NumAccess, "Wrong! NumAccess should be larger than 0");
static_assert(NumAccess % NumCoord == 0, "wrong! # of access is not divisible by NumCoord");
};
static constexpr index_t NumAccessPerCoord = load_store_traits::NumAccess / NumCoord;
CK_TILE_DEVICE constexpr tile_scatter_gather() = default;
CK_TILE_DEVICE constexpr tile_scatter_gather(
const BottomTensorView& bottom_tensor_view,
const WindowLengths& window_lengths,
const BottomTensorIndex& window_origin,
const TileDstr& tile_distribution,
const PageIdxArray& page_idx)
: bottom_tensor_view_{bottom_tensor_view},
window_lengths_{window_lengths},
window_origin_{window_origin},
tile_dstr_{tile_distribution},
page_idx_{page_idx},
pre_computed_coords_{}
{
#if 0 // debug
// TODO: this use more register for FA, but less register for GEMM
// need investigation
// only support warp-tile and block-tile
static_assert(NDimP == 1 or NDimP == 2, "wrong!");
WindowAdaptorCoord window_adaptor_thread_coord_tmp;
if constexpr(NDimP == 1)
{
window_adaptor_thread_coord_tmp = make_tensor_adaptor_coordinate(
tile_distribution.get_ps_ys_to_xs_adaptor(), AdaptorTopIndex{get_lane_id(), 0});
}
else if constexpr(NDimP == 2)
{
window_adaptor_thread_coord_tmp =
make_tensor_adaptor_coordinate(tile_distribution.get_ps_ys_to_xs_adaptor(),
AdaptorTopIndex{get_warp_id(), get_lane_id(), 0});
}
#else
// TODO: this use less register for FA, but more register for GEMM
// need investigation
const auto window_adaptor_thread_coord_tmp = make_tensor_adaptor_coordinate(
tile_distribution.get_ps_ys_to_xs_adaptor(),
container_concat(detail::get_partition_index(tile_distribution),
array<index_t, NDimY>{0}));
#endif
BottomTensorIndex bottom_tensor_thread_origin_idx_tmp =
window_origin + window_adaptor_thread_coord_tmp.get_bottom_index();
bottom_tensor_thread_origin_idx_tmp(HsGatherDim) = 0;
const auto bottom_tensor_thread_coord_tmp = make_tensor_coordinate(
bottom_tensor_view_.get_tensor_descriptor(), bottom_tensor_thread_origin_idx_tmp);
// pre-compute NumCoord (WindowAdaptorCoord, BottomTensorCoord) bundles to speed up
// future load/store() calls (might allocate more registers)
using Traits = load_store_traits;
using SFC_Ys = typename Traits::SFC_Ys;
static_for<0, NumCoord, 1>{}([&](auto iCoord) {
auto window_adaptor_thread_coord = window_adaptor_thread_coord_tmp;
auto bottom_tensor_thread_coord = bottom_tensor_thread_coord_tmp;
constexpr auto idx_diff_ys =
SFC_Ys::get_step_between(number<0>{}, number<iCoord * NumAccessPerCoord>{});
constexpr auto idx_diff_ps_ys = container_concat(
generate_tuple([&](auto) { return number<0>{}; }, number<NDimP>{}), idx_diff_ys);
move_window_adaptor_and_bottom_tensor_thread_coordinate(
window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys);
pre_computed_coords_(iCoord) =
make_tuple(window_adaptor_thread_coord, bottom_tensor_thread_coord);
});
}
CK_TILE_DEVICE static constexpr index_t get_num_of_dimension() { return NDimBottomTensor; }
CK_TILE_DEVICE static constexpr bool has_static_tile_distribution()
{
return TileDstr::is_static();
}
CK_TILE_DEVICE constexpr auto get_window_lengths() const { return window_lengths_; }
CK_TILE_DEVICE constexpr auto get_tile_distribution() const { return tile_dstr_; }
CK_TILE_DEVICE constexpr auto get_bottom_tensor_view() const { return bottom_tensor_view_; }
CK_TILE_DEVICE constexpr auto get_window_origin() const { return window_origin_; }
CK_TILE_DEVICE constexpr void
set_bottom_tensor_view_data_ptr(typename BottomTensorView::DataType* data)
{
bottom_tensor_view_.buf_.p_data_ = data;
}
// move thread's window adaptor coordinate and bottom tensor coordinate
// [p0, p1, ..., y0, y1, ...] ==> [x0, x1, ...] ==> [x0', x1', ...] ==> [offset]
template <typename ATopIndex>
CK_TILE_DEVICE void move_window_adaptor_and_bottom_tensor_thread_coordinate(
WindowAdaptorCoord& window_adaptor_thread_coord,
BottomTensorCoord& bottom_tensor_thread_coord,
const ATopIndex& idx_diff_adaptor_top) const
{
array<index_t, NDimBottomTensor> idx_diff_adaptor_bottom;
move_tensor_adaptor_coordinate(tile_dstr_.get_ps_ys_to_xs_adaptor(),
window_adaptor_thread_coord,
idx_diff_adaptor_top,
idx_diff_adaptor_bottom);
move_tensor_coordinate(bottom_tensor_view_.get_tensor_descriptor(),
bottom_tensor_thread_coord,
idx_diff_adaptor_bottom);
}
// return vector dimension among [y0, y1, ...]
CK_TILE_DEVICE static constexpr auto get_window_adaptor_ys_safe_vector_length_strides()
{
// bottom tensor top dimension vector lengths and strides
const auto [bottom_tensor_top_dim_vector_lengths, bottom_tensor_top_dim_vector_strides] =
BottomTensorDesc::get_top_dimension_safe_vector_length_strides();
// window vector lengths/strides
const auto window_adaptor_bottom_dim_vector_lengths = bottom_tensor_top_dim_vector_lengths;
const auto window_adaptor_bottom_dim_vector_strides = bottom_tensor_top_dim_vector_strides;
// window adaptor [p0, p1, ..., y0, y1, ...]
array<index_t, WindowAdaptor::get_num_of_hidden_dimension()> window_adaptor_vector_lengths{
-1};
array<index_t, WindowAdaptor::get_num_of_hidden_dimension()> window_adaptor_vector_strides{
-1};
constexpr auto window_adaptor_bottom_dims =
WindowAdaptor::get_bottom_dimension_hidden_ids();
set_container_subset(window_adaptor_vector_lengths,
window_adaptor_bottom_dims,
window_adaptor_bottom_dim_vector_lengths);
set_container_subset(window_adaptor_vector_strides,
window_adaptor_bottom_dims,
window_adaptor_bottom_dim_vector_strides);
const auto [window_adaptor_ps_ys_vector_lengths, window_adaptor_ps_ys_vector_strides] =
WindowAdaptor{}.get_top_dimension_safe_vector_length_strides(
window_adaptor_vector_lengths, window_adaptor_vector_strides);
// [y0, y1, ...]
constexpr auto y_dims = typename arithmetic_sequence_gen<TileDstr::get_num_of_dimension_p(),
NDimWindowAdaptorTop,
1>::type{};
return make_tuple(get_container_subset(window_adaptor_ps_ys_vector_lengths, y_dims),
get_container_subset(window_adaptor_ps_ys_vector_strides, y_dims));
}
CK_TILE_DEVICE constexpr auto get_num_of_access() const { return load_store_traits::NumAccess; }
template <index_t i_access_unsupport_ = -1, bool oob_conditional_check = true>
CK_TILE_DEVICE auto load(number<i_access_unsupport_> = {},
bool_constant<oob_conditional_check> = {}) const
{
constexpr auto tile_dstr = TileDstr{};
auto dst_tensor = make_static_distributed_tensor<DataType>(tile_dstr);
load(dst_tensor, number<i_access_unsupport_>{}, bool_constant<oob_conditional_check>{});
return dst_tensor;
}
template <typename DistributedTensor,
index_t i_access_unsupport_ = -1,
bool oob_conditional_check = true>
CK_TILE_DEVICE auto load(DistributedTensor& dst_tensor,
number<i_access_unsupport_> = {},
bool_constant<oob_conditional_check> = {}) const
{
using Traits = load_store_traits;
using vector_t = typename Traits::vector_t;
using SFC_Ys = typename Traits::SFC_Ys;
constexpr auto tile_dstr = TileDstr{};
// loop over thread tensor space [y0, y1, ...]
static_for<0, NumCoord, 1>{}([&](auto iCoord) {
/// TODO: use structure binding (to be captured later) if compiled in C++20
auto window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0];
auto bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1];
static_for<0, NumAccessPerCoord, 1>{}([&](auto iCoordAccess) {
constexpr auto iAccess = number<iCoord * NumAccessPerCoord + iCoordAccess>{};
// data index [y0, y1, ...]
constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess);
constexpr auto idx_gather = idx_ys_start[number<YsGatherDim>{}];
const auto page_offset = page_idx_[idx_gather];
// read from bottom tensor
const vector_t vec_value =
get_bottom_tensor_view().template get_vectorized_elements<vector_t>(
bottom_tensor_thread_coord, page_offset, bool_constant<oob_conditional_check>{});
#if 1
// write into distributed tensor
static_for<0, Traits::ScalarPerVector, Traits::PackedSize>{}([&](auto j) {
constexpr auto idx_ys = generate_tuple(
[&](auto jj) {
return jj == Traits::VectorDimY ? (idx_ys_start[jj] + j)
: idx_ys_start[jj];
},
number<NDimY>{});
constexpr index_t d =
tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys) /
Traits::PackedSize;
dst_tensor.get_thread_buffer().template at<d>() =
vec_value.template get_as<DataType>()[j / Traits::PackedSize];
});
#else
constexpr index_t d =
tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys_start);
static_assert(d % Traits::ScalarPerVector == 0);
dst_tensor.get_thread_buffer().template get_as<vector_t>()(
number<d / Traits::ScalarPerVector>{}) = bit_cast<vector_t>(vec_value);
#endif
// move thread coordinate
if constexpr(iCoordAccess != (NumAccessPerCoord - 1))
{
constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess);
constexpr auto forward_step_scatter = generate_tuple(
[&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, number<NDimY>{});
constexpr auto idx_diff_ps_ys = container_concat(
generate_tuple([&](auto) { return number<0>{}; }, number<NDimP>{}),
forward_step_scatter);
move_window_adaptor_and_bottom_tensor_thread_coordinate(
window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys);
}
});
});
}
// TODO: currently async load only implemented in inline asm
template <typename LdsTileWindow_,
index_t i_access_unsupport_ = -1,
bool oob_conditional_check = true,
bool pre_nop = false>
CK_TILE_DEVICE auto async_load_raw(LdsTileWindow_&& lds_tile,
number<i_access_unsupport_> = {},
bool_constant<oob_conditional_check> = {},
bool_constant<pre_nop> = {}) const
{
using LdsTileWindow = remove_cvref_t<LdsTileWindow_>;
// using LdsTensorView = typename LdsTileWindow::BottomTensorView;
using LdsDataType = typename LdsTileWindow::DataType;
// using LdsDescriptor = typename LdsTileWindow::BottomTensorDesc;
// issues * warps * lanes
static_assert(LdsTileWindow::get_num_of_dimension() == 3); // TODO: hard coded
const index_t size_per_buf =
lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset(
make_tuple(number<0>{}, number<0>{}, number<0>{})) *
sizeof(LdsDataType);
const index_t size_per_wave =
lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset(
make_tuple(number<0>{}, number<1>{}, number<0>{})) *
sizeof(LdsDataType) -
size_per_buf;
const index_t size_per_issue =
lds_tile.get_bottom_tensor_view().get_tensor_descriptor().calculate_offset(
make_tuple(number<1>{}, number<0>{}, number<0>{})) *
sizeof(LdsDataType) -
size_per_buf;
const index_t m0_init_value = size_per_buf + size_per_wave * get_warp_id();
m0_set_with_memory(m0_init_value); // This should be wave independent
using Traits = load_store_traits;
// using vector_type_t = typename Traits::vector_type_t;
using vector_t = typename Traits::vector_t;
using SFC_Ys = typename Traits::SFC_Ys;
LdsDataType* smem = lds_tile.get_bottom_tensor_view().get_buffer_view().p_data_;
// loop over thread tensor space [y0, y1, ...]
static_for<0, NumCoord, 1>{}([&](auto iCoord) {
/// TODO: use structure binding (to be captured later) if compiled in C++20
auto window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0];
auto bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1];
static_for<0, NumAccessPerCoord, 1>{}([&](auto iCoordAccess) {
constexpr auto iAccess = number<iCoord * NumAccessPerCoord + iCoordAccess>{};
constexpr auto pre_nop_ = [&]() {
if constexpr(pre_nop && iCoord == 0 && iCoordAccess == 0)
return bool_constant<true>{};
else
return bool_constant<false>{};
}();
constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess);
constexpr auto idx_gather = idx_ys_start[number<YsGatherDim>{}];
const auto page_offset = page_idx_[idx_gather];
// read from bottom tensor
get_bottom_tensor_view().template async_get_vectorized_elements_raw<vector_t>(
smem, bottom_tensor_thread_coord, page_offset, 0, pre_nop_);
// move thread coordinate
if constexpr(iCoordAccess != (NumAccessPerCoord - 1))
{
constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess);
constexpr auto forward_step_scatter = generate_tuple(
[&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, number<NDimY>{});
constexpr auto idx_diff_ps_ys = container_concat(
generate_tuple([&](auto) { return number<0>{}; }, number<NDimP>{}),
forward_step_scatter);
move_window_adaptor_and_bottom_tensor_thread_coordinate(
window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys);
m0_inc_with_memory(size_per_issue);
}
});
});
}
template <index_t i_access_unsupport_ = -1,
bool oob_conditional_check = true>
CK_TILE_DEVICE void store(const static_distributed_tensor<DataType, TileDstr>& dstr_tensor,
number<i_access_unsupport_> = {},
bool_constant<oob_conditional_check> = {}) const
{
using Traits = load_store_traits;
// using vector_type_t = typename Traits::vector_type_t;
using vector_t = typename Traits::vector_t;
using SFC_Ys = typename Traits::SFC_Ys;
constexpr auto tile_dstr = TileDstr{};
// printf("off %d\n", page_idx_[I0]);
// loop over thread tensor space [y0, y1, ...]
static_for<0, NumCoord, 1>{}([&](auto iCoord) {
auto window_adaptor_thread_coord = pre_computed_coords_[iCoord][I0];
auto bottom_tensor_thread_coord = pre_computed_coords_[iCoord][I1];
static_for<0, NumAccessPerCoord, 1>{}([&](auto iCoordAccess) {
constexpr auto iAccess = number<iCoord * NumAccessPerCoord + iCoordAccess>{};
// data index [y0, y1, ...]
constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess);
constexpr auto idx_gather = idx_ys_start[number<0>{}];
const auto page_offset = page_idx_[idx_gather];
// printf("idx_ys_start[0], idx_ys_start[1](%d, %d) \n",
// idx_ys_start[number<0>{}]+0, idx_ys_start[number<1>{}]+0);
// read from distributed tensor
// vector_type_t vec;
vector_t vec_value;
static_for<0, Traits::ScalarPerVector, Traits::PackedSize>{}([&](auto j) {
constexpr auto idx_ys = generate_tuple(
[&](auto jj) {
return jj == Traits::VectorDimY ? (idx_ys_start[jj] + j)
: idx_ys_start[jj];
},
number<NDimY>{});
constexpr index_t d =
tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys) /
Traits::PackedSize;
// printf("thread_idx_m: %d j: %d\n", idx_ys[number<0>{}] + 0, 0+j);
vec_value.template get_as<DataType>()(j / Traits::PackedSize) =
dstr_tensor.get_thread_buffer().template at<d>();
});
// const vector_t vec_value = vec.template get_as<vector_t>().template at<0>();
// write into bottom tensor
get_bottom_tensor_view().template set_vectorized_elements<vector_t>(
bottom_tensor_thread_coord,
page_offset,
vec_value,
bool_constant<oob_conditional_check>{});
// printf("coord_offset:%d, scatter_offset:%d \n",
// bottom_tensor_thread_coord.get_offset(), offset); move thread coordinate
if constexpr(iCoordAccess != (NumAccessPerCoord - 1))
{
constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess);
constexpr auto forward_step_scatter = generate_tuple(
[&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, number<NDimY>{});
constexpr auto idx_diff_ps_ys = container_concat(
generate_tuple([&](auto) { return number<0>{}; }, number<NDimP>{}),
forward_step_scatter);
move_window_adaptor_and_bottom_tensor_thread_coordinate(
window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys);
}
});
});
}
// move thread's botom tensor coordiante
// [x0', x1', ... ] ==> [offset]
// also move window-origin
CK_TILE_DEVICE void move(const BottomTensorIndex& step)
{
window_origin_ += step;
BottomTensorIndex step_new = step;
step_new(HsGatherDim) = 0;
static_for<0, NumCoord, 1>{}([&](auto iCoord) {
move_tensor_coordinate(bottom_tensor_view_.get_tensor_descriptor(),
pre_computed_coords_(iCoord)(I1),
step_new);
});
}
CK_TILE_DEVICE void update_page_idx(const PageIdxArray& new_idx)
{
page_idx_ = new_idx;
// static_for<0, 2, 1>{}([&](auto k0) {
// printf("update tid %d %d \n", threadIdx.x, page_idx_[k0]);
// });
}
CK_TILE_DEVICE void set_window_origin(const BottomTensorIndex& new_window_origin)
{
window_origin_ = new_window_origin;
#if 0 // debug
// TODO: this use more register for FA, but less register for GEMM
// need investigation
// only support warp-tile and block-tile
static_assert(NDimP == 1 or NDimP == 2, "wrong!");
WindowAdaptorCoord window_adaptor_thread_coord_tmp;
if constexpr(NDimP == 1)
{
window_adaptor_thread_coord_tmp = make_tensor_adaptor_coordinate(
tile_dstr_.get_ps_ys_to_xs_adaptor(), AdaptorTopIndex{get_lane_id(), 0});
}
else if constexpr(NDimP == 2)
{
window_adaptor_thread_coord_tmp =
make_tensor_adaptor_coordinate(tile_dstr_.get_ps_ys_to_xs_adaptor(),
AdaptorTopIndex{get_warp_id(), get_lane_id(), 0});
}
#else
// TODO: this use less register for FA, but more register for GEMM
// need investigation
const auto window_adaptor_thread_coord_tmp = make_tensor_adaptor_coordinate(
tile_dstr_.get_ps_ys_to_xs_adaptor(),
container_concat(detail::get_partition_index(tile_dstr_), array<index_t, NDimY>{0}));
#endif
BottomTensorIndex bottom_tensor_thread_origin_idx_tmp =
window_origin_ + window_adaptor_thread_coord_tmp.get_bottom_index();
bottom_tensor_thread_origin_idx_tmp(HsGatherDim) = 0;
const auto bottom_tensor_thread_coord_tmp = make_tensor_coordinate(
bottom_tensor_view_.get_tensor_descriptor(), bottom_tensor_thread_origin_idx_tmp);
// pre-compute NumCoord (WindowAdaptorCoord, BottomTensorCoord) bundles to speed up
// future load/store() calls (might allocate more registers)
using Traits = load_store_traits;
using SFC_Ys = typename Traits::SFC_Ys;
static_for<0, NumCoord, 1>{}([&](auto iCoord) {
auto window_adaptor_thread_coord = window_adaptor_thread_coord_tmp;
auto bottom_tensor_thread_coord = bottom_tensor_thread_coord_tmp;
constexpr auto idx_diff_ys =
SFC_Ys::get_step_between(number<0>{}, number<iCoord * NumAccessPerCoord>{});
constexpr auto idx_diff_ps_ys = container_concat(
generate_tuple([&](auto) { return number<0>{}; }, number<NDimP>{}), idx_diff_ys);
move_window_adaptor_and_bottom_tensor_thread_coordinate(
window_adaptor_thread_coord, bottom_tensor_thread_coord, idx_diff_ps_ys);
pre_computed_coords_(iCoord) =
make_tuple(window_adaptor_thread_coord, bottom_tensor_thread_coord);
});
}
CK_TILE_HOST_DEVICE void init_raw() { bottom_tensor_view_.init_raw(); }
// this is the bottom tensor view
// [x0', x1', ...] ==> [offset]
BottomTensorView bottom_tensor_view_;
//
WindowLengths window_lengths_;
// origin ([x0', x1', ...]) of window on bottom tensor
BottomTensorIndex window_origin_;
// Tile tensor distribution, which contains:
// 1. adaptor for window: [p0, p1, ..., y0, y1, ...] ==> [x0, x1, ...]
// 2. thread descriptor for thread tensor in register: [y0, y1, ...] ==> [d]
TileDstr tile_dstr_;
PageIdxArray page_idx_;
// this contains:
// per-thread coordinate for window adaptor
// per-thread coordinate for bottom tensor
array<tuple<WindowAdaptorCoord, BottomTensorCoord>, NumCoord> pre_computed_coords_;
};
// TODO: use strategy
template <typename TensorView_,
typename WindowLengths_,
typename StaticTileDistribution_,
typename StaticPageIndexArray_,
index_t HsGatherDim = 0,
index_t NumCoord = 1>
CK_TILE_DEVICE constexpr auto
make_tile_scatter_gather(const TensorView_& tensor_view,
const WindowLengths_& window_lengths,
const multi_index<TensorView_::get_num_of_dimension()>& origin,
const StaticTileDistribution_& tile_distribution,
const StaticPageIndexArray_& page_idx,
number<HsGatherDim> = {},
number<NumCoord> = {})
{
return tile_scatter_gather<remove_cvref_t<TensorView_>,
remove_cvref_t<WindowLengths_>,
remove_cvref_t<StaticTileDistribution_>,
remove_cvref_t<StaticPageIndexArray_>,
HsGatherDim,
NumCoord>{
tensor_view, window_lengths, origin, tile_distribution, page_idx};
}
template <typename TensorView, typename WindowLengths, typename StaticTileDistribution, typename StaticPageIndexArray, index_t HsGatherDim>
CK_TILE_DEVICE constexpr auto
make_tile_scatter_gather(const tile_window_with_static_lengths<TensorView, WindowLengths>& tile_window,
const multi_index<TensorView::get_num_of_dimension()>& origin,
const StaticTileDistribution& tile_distribution,
const StaticPageIndexArray& page_idx,
number<HsGatherDim> = {})
{
return make_tile_scatter_gather(tile_window.get_bottom_tensor_view(),
tile_window.get_window_lengths(),
origin,
tile_distribution,
page_idx,
number<HsGatherDim>{});
}
template <typename TensorView, typename WindowLengths, typename StaticTileDistribution, typename StaticPageIndexArray, index_t HsGatherDim>
CK_TILE_DEVICE constexpr auto
make_tile_scatter_gather(const tile_window_with_static_lengths<TensorView, WindowLengths>& tile_window,
const StaticTileDistribution& tile_distribution, const StaticPageIndexArray& page_idx,
number<HsGatherDim> = {})
{
return make_tile_scatter_gather(tile_window.get_bottom_tensor_view(),
tile_window.get_window_lengths(),
tile_window.get_window_origin(),
tile_distribution,
page_idx,
number<HsGatherDim>{});
}
// template <typename TensorView, typename WindowLengths, typename StaticTileDistribution>
// CK_TILE_DEVICE constexpr auto
// make_tile_window_raw(const tile_window_with_static_lengths<TensorView, WindowLengths>& tile_window,
// const StaticTileDistribution& tile_distribution)
// {
// auto w = make_tile_scatter_gather(tile_window.get_bottom_tensor_view(),
// tile_window.get_window_lengths(),
// tile_window.get_window_origin(),
// tile_distribution);
// w.init_raw();
// return w;
// }
} // namespace ck_tile

28
include/ck_tile/core/tensor/tile_window.hpp
View File

@@ -613,7 +613,7 @@ struct tile_window_with_static_distribution
index_t i_access_unsupport_ = -1,
bool oob_conditional_check = true>
CK_TILE_DEVICE void store(const static_distributed_tensor<DataType, TileDstr>& dstr_tensor,
const statically_indexed_array offsets,
const statically_indexed_array& offsets,
number<i_access_unsupport_> = {},
bool_constant<oob_conditional_check> = {}) const
{
@@ -1097,23 +1097,6 @@ make_tile_window_raw(const TensorView_& tensor_view,
return w;
}
template <typename TensorView_,
typename WindowLengths_,
typename StaticTileDistribution_,
index_t NumCoord>
CK_TILE_DEVICE void move_tile_window(
tile_window_with_static_distribution<TensorView_,
WindowLengths_,
StaticTileDistribution_,
NumCoord>& window,
const typename tile_window_with_static_distribution<TensorView_,
WindowLengths_,
StaticTileDistribution_,
NumCoord>::BottomTensorIndex& step)
{
window.move(step);
}
/**
* @brief This class provides description of tile windowed view on the device memory.
*
@@ -1242,13 +1225,4 @@ make_tile_window_raw(const tile_window_with_static_lengths<TensorView, WindowLen
return w;
}
template <typename TensorView_, typename WindowLengths_>
CK_TILE_DEVICE void move_tile_window(
tile_window_with_static_lengths<TensorView_, WindowLengths_>& window,
const typename tile_window_with_static_lengths<TensorView_, WindowLengths_>::BottomTensorIndex&
step)
{
window.move(step);
}
} // namespace ck_tile

15
include/ck_tile/core/tensor/tile_window_linear.hpp
View File

@@ -1200,19 +1200,4 @@ make_tile_window_linear_raw(const TileWindow_& tile_window,
LinearBottomDims_{});
}
template <typename TensorView_,
typename WindowLengths_,
typename StaticTileDistribution_,
typename LinearBottomDims_>
CK_TILE_DEVICE void move_tile_window(
tile_window_linear<TensorView_, WindowLengths_, StaticTileDistribution_, LinearBottomDims_>&
window,
const typename tile_window_linear<TensorView_,
WindowLengths_,
StaticTileDistribution_,
LinearBottomDims_>::BottomTensorIndex& step)
{
window.move(step);
}
} // namespace ck_tile

8
include/ck_tile/core/tensor/tile_window_utils.hpp
View File

@@ -18,6 +18,14 @@
#pragma once
namespace ck_tile {
template <typename TileWindow_>
CK_TILE_DEVICE void move_tile_window(
TileWindow_& window,
const typename TileWindow_::BottomTensorIndex& step)
{
window.move(step);
}
// input a lds store tile, extract some information from it
// used to set m0 value for gfx9 serious
template <typename LdsTileWindow_>

108
include/ck_tile/core/tensor/transpose_tile.hpp
View File

@@ -83,6 +83,9 @@ CK_TILE_DEVICE void transpose_tile2d_impl_in_thread(OutTensor& out_tensor,
constexpr index_t num_vec_in = vec_length_out;
constexpr index_t num_vec_out = vec_length_in;
using InVec = array<DataType, vec_length_in>;
using OutVec = array<DataType, vec_length_out>;
// SFC
constexpr auto scalars_per_access_arr = generate_array(
[&](auto i) { return (i == y_dim_vec_in or i == y_dim_vec_out) ? y_lengths[i] : 1; },
@@ -98,84 +101,51 @@ CK_TILE_DEVICE void transpose_tile2d_impl_in_thread(OutTensor& out_tensor,
static_assert(num_access > 0, "wrong! num_access should be larger than 0");
if constexpr(num_vec_in == 1 || num_vec_out == 1)
{
// loop over SFC
static_for<0, num_access, 1>{}([&](auto iAccess) {
// data index [y0, y1, ...] in the order of input tensor
constexpr auto idx_y = SFC_Y::get_index(iAccess);
// in/out vectors to be transposed
thread_buffer<InVec, num_vec_in> in_vectors;
thread_buffer<OutVec, num_vec_out> out_vectors;
constexpr index_t in_offset = y_in_desc.calculate_offset(idx_y);
constexpr index_t out_offset = y_out_desc.calculate_offset(idx_y);
// loop over SFC and do transpose
static_for<0, num_access, 1>{}([&](auto iAccess) {
// data index [y0, y1, ...] in the order of input tensor
constexpr auto idx_y_start = SFC_Y::get_index(iAccess);
if constexpr(vec_length_in == 1)
{
out_tensor.get_thread_buffer()[number<out_offset>{}] =
in_tensor.get_thread_buffer()[number<in_offset>{}];
}
else
{
using Vec = array<DataType, vec_length_in>;
out_tensor.get_thread_buffer().template get_as<Vec>(
number<out_offset / vec_length_in>{}) =
in_tensor.get_thread_buffer().template get_as<Vec>(
number<in_offset / vec_length_in>{});
}
// get input vectors
static_for<0, num_vec_in, 1>{}([&](auto i) {
constexpr auto idx_y_in = generate_tuple(
[&](auto ii) {
return ii == y_dim_vec_out ? idx_y_start[ii] + i : idx_y_start[ii];
},
number<NDimY>{});
constexpr index_t in_offset = y_in_desc.calculate_offset(idx_y_in);
static_assert(in_offset % vec_length_in == 0);
in_vectors(i).template get_as<InVec>()(I0) =
in_tensor.get_thread_buffer()
.template get_as<InVec>()[number<in_offset / vec_length_in>{}];
});
}
else
{
using InVec = array<DataType, vec_length_in>;
using OutVec = array<DataType, vec_length_out>;
// in/out vectors to be transposed
thread_buffer<InVec, num_vec_in> in_vectors;
thread_buffer<OutVec, num_vec_out> out_vectors;
// transpose
transpose_vectors<DataType, num_vec_in, num_vec_out>{}(in_vectors, out_vectors);
// loop over SFC and do transpose
static_for<0, num_access, 1>{}([&](auto iAccess) {
// data index [y0, y1, ...] in the order of input tensor
constexpr auto idx_y_start = SFC_Y::get_index(iAccess);
// set output vectors
static_for<0, num_vec_out, 1>{}([&](auto i) {
constexpr auto idx_y_out_tmp = generate_array(
[&](auto ii) { return ii == y_dim_vec_in ? idx_y_start[ii] + i : idx_y_start[ii]; },
number<NDimY>{});
// get input vectors
static_for<0, num_vec_in, 1>{}([&](auto i) {
constexpr auto idx_y_in = generate_tuple(
[&](auto ii) {
return ii == y_dim_vec_out ? idx_y_start[ii] + i : idx_y_start[ii];
},
number<NDimY>{});
constexpr auto idx_y_out =
container_reorder_given_new2old(idx_y_out_tmp, y_dim_out_to_in);
constexpr index_t in_offset = y_in_desc.calculate_offset(idx_y_in);
static_assert(in_offset % vec_length_in == 0);
constexpr index_t out_offset = y_out_desc.calculate_offset(idx_y_out);
static_assert(out_offset % vec_length_out == 0);
in_vectors(i).template get_as<InVec>()(I0) =
in_tensor.get_thread_buffer()
.template get_as<InVec>()[number<in_offset / vec_length_in>{}];
});
// transpose
transpose_vectors<DataType, num_vec_in, num_vec_out>{}(in_vectors, out_vectors);
// set output vectors
static_for<0, num_vec_out, 1>{}([&](auto i) {
constexpr auto idx_y_out_tmp = generate_array(
[&](auto ii) {
return ii == y_dim_vec_in ? idx_y_start[ii] + i : idx_y_start[ii];
},
number<NDimY>{});
constexpr auto idx_y_out =
container_reorder_given_new2old(idx_y_out_tmp, y_dim_out_to_in);
constexpr index_t out_offset = y_out_desc.calculate_offset(idx_y_out);
static_assert(out_offset % vec_length_out == 0);
out_tensor.get_thread_buffer().template set_as<OutVec>(
number<out_offset / vec_length_out>{},
out_vectors[i].template get_as<OutVec>()[I0]);
});
out_tensor.get_thread_buffer().template set_as<OutVec>(
number<out_offset / vec_length_out>{},
out_vectors[i].template get_as<OutVec>()[I0]);
});
}
});
}
} // namespace detail

12
include/ck_tile/host/check_err.hpp
View File

@@ -186,7 +186,7 @@ check_err(const Range& out,
{
max_err = err > max_err ? err : max_err;
err_count++;
if(err_count < 5)
if(err_count < 5000000)
{
std::cerr << msg << std::setw(12) << std::setprecision(7) << " out[" << i
<< "] != ref[" << i << "]: " << o << " != " << r << std::endl;
@@ -246,7 +246,7 @@ check_err(const Range& out,
{
max_err = err > max_err ? err : max_err;
err_count++;
if(err_count < 5)
if(err_count < 5000000)
{
std::cerr << msg << std::setw(12) << std::setprecision(7) << " out[" << i
<< "] != ref[" << i << "]: " << o << " != " << r << std::endl;
@@ -305,7 +305,7 @@ check_err(const Range& out,
{
max_err = err > max_err ? err : max_err;
err_count++;
if(err_count < 5)
if(err_count < 5000000)
{
std::cerr << msg << std::setw(12) << std::setprecision(7) << " out[" << i
<< "] != ref[" << i << "]: " << o << " != " << r << std::endl;
@@ -360,7 +360,7 @@ std::enable_if_t<(std::is_same_v<ranges::range_value_t<Range>, ranges::range_val
{
max_err = err > max_err ? err : max_err;
err_count++;
if(err_count < 5)
if(err_count < 5000000)
{
std::cerr << msg << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r
<< std::endl;
@@ -437,7 +437,7 @@ std::enable_if_t<(std::is_same_v<ranges::range_value_t<Range>, ranges::range_val
{
max_err = err > max_err ? err : max_err;
err_count++;
if(err_count < 5)
if(err_count < 5000000)
{
std::cerr << msg << std::setw(12) << std::setprecision(7) << " out[" << i
<< "] != ref[" << i << "]: " << o_fp64 << " != " << r_fp64 << std::endl;
@@ -495,7 +495,7 @@ std::enable_if_t<(std::is_same_v<ranges::range_value_t<Range>, ranges::range_val
{
max_err = err > max_err ? err : max_err;
err_count++;
if(err_count < 5)
if(err_count < 5000000)
{
std::cerr << msg << std::setw(12) << std::setprecision(7) << " out[" << i
<< "] != ref[" << i << "]: " << o << " != " << r << std::endl;

1
include/ck_tile/host/reference/reference_fused_single_moe_gemm.hpp
View File

@@ -98,6 +98,7 @@ __global__ void naive_gemm_kernel(const ck_tile::index_t* p_sorted_token_ids_,
int row = idx / N; // Compute row index
int col = idx % N; // Compute column index
(void)Num_tokens;
// assert(p_sorted_expert_ids_ != nullptr);
// assert(TopK == 1);
// assert(Num_tokens == 128);

41
include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp
View File

@@ -119,6 +119,17 @@ struct CShuffleEpilogue
return kMWave * kNWave * kMPerXdl * kNPerXdl * sizeof(ODataType);
}
CK_TILE_HOST_DEVICE static constexpr auto GetCDisrtribution()
{
using TileEncodingPattern =
TileDistributionEncodingPattern2D<kBlockSize,
kMPerIteration,
kNPerIteration,
GetVectorSizeC(),
tile_distribution_pattern::thread_raked>;
return TileEncodingPattern::Make2DStaticTileDistribution();
}
template <typename ODramWindow,
typename OAccTile,
bool IsInputGemm = true,
@@ -127,7 +138,9 @@ struct CShuffleEpilogue
const OAccTile& o_acc_tile,
void* p_smem,
const index_t* p_sorted_tokens_id,
index_t token_pos)
index_t token_pos,
index_t TopK,
index_t stride_C)
{
const index_t iMWarp = get_warp_id() / kNWave;
@@ -174,21 +187,25 @@ struct CShuffleEpilogue
// auto idx_m = number<idx_y_start.at(number<0>{})>{} + 0;
// printf("idx_y_start:%d \n", idx_m);
constexpr auto mIter = number<idx_y_start.at(number<0>{}) / (kMPerXdl * kMWave)>{};
constexpr auto MPerAcess = kMPerXdl * kMWave;
constexpr auto mIter = number<idx_y_start.at(number<0>{}) / (MPerAcess)>{};
using CDstrEncode = typename decltype(dram_tile_distribution)::DstrEncode;
constexpr ck_tile::index_t MRepeat = CDstrEncode::hs_lengthss_[number<0>{}][number<0>{}];
statically_indexed_array<index_t, 2> offsets;
static_for<0, 2 /*CMrepeats*/, 1>{}([&](auto m0) {
statically_indexed_array<index_t, MRepeat> offsets;
static_for<0, MRepeat, 1>{}([&](auto m0) {
auto token_id = token_pos + m0 + c_coord[0] + mIter * kMPerXdl * kMWave;
// auto token_id = token_pos + c_coord[0] + mIter * MPerAcess + MPerAcess / MRepeat * m0.value;
auto fused_token = p_sorted_tokens_id[token_id];
index_t token_offset = fused_token & 0xffffff;
index_t scatter_token_id = fused_token & 0xffffff;
if constexpr(IsInputGemm)
{
token_offset = token_offset * 3 /*TopK*/ + (fused_token >> 24);
scatter_token_id = scatter_token_id * TopK + (fused_token >> 24);
}
offsets[m0] = token_offset * 4096; // Problem::kN_;
offsets[m0] = scatter_token_id * stride_C; // Problem::kN_;
});
// printf("c_coord[number<0>{}]: %d \n", coord[number<0>{}]);
// printf("mIter: %d", mIter+0);
@@ -212,7 +229,13 @@ struct CShuffleEpilogue
if constexpr(out_memory_data_op == memory_operation_enum::set)
{
store_tile(out_dram_window, c_out_tensor, offsets);
auto tile_window = make_tile_scatter_gather(out_dram_window.get_bottom_tensor_view(),
out_dram_window.get_window_lengths(),
out_dram_window.get_window_origin(),
dram_tile_distribution,
offsets);
tile_window.store(c_out_tensor);
// store_tile(out_dram_window, c_out_tensor, offsets);
}
else
{

2
include/ck_tile/ops/gemm.hpp
View File

@@ -26,7 +26,6 @@
#include "ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp"
#include "ck_tile/ops/gemm/kernel/gemm_kernel.hpp"
#include "ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp"
#include "ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp"
#include "ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp"
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp"
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp"
@@ -40,7 +39,6 @@
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2_default_policy.hpp"
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp"
#include "ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp"
#include "ck_tile/ops/gemm/pipeline//moe_gemm_pipeline_ag_bg_cr.hpp"
#include "ck_tile/ops/gemm/pipeline/tile_gemm_shape.hpp"
#include "ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp"
#include "ck_tile/ops/gemm/warp/warp_gemm.hpp"

506
include/ck_tile/ops/gemm/kernel/moe_gemm_kernel.hpp
View File

@@ -1,506 +0,0 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck_tile/core/numeric/math.hpp"
#include "ck_tile/core/utility/literals.hpp"
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp"
#include "ck_tile/ops/gemm/kernel/gemm_kernel.hpp"
#include "ck_tile/host.hpp"
namespace ck_tile {
struct MoeGemmHostArgs : public ck_tile::GemmHostArgs
{
ck_tile::index_t NumTokens;
ck_tile::index_t TopK;
const ck_tile::index_t* p_sorted_token_ids;
const ck_tile::index_t* p_sorted_expert_ids;
const ck_tile::index_t* p_max_token_id;
// TODO: add kbatch for splitk
CK_TILE_HOST MoeGemmHostArgs() noexcept = default;
CK_TILE_HOST MoeGemmHostArgs(const ck_tile::index_t* p_sorted_token_ids_,
const ck_tile::index_t* p_sorted_expert_ids_,
const ck_tile::index_t* p_max_token_id_,
const void* a_ptr_,
const void* b_ptr_,
void* c_ptr_,
ck_tile::index_t NumTokens_,
ck_tile::index_t TopK_,
ck_tile::index_t M_,
ck_tile::index_t N_,
ck_tile::index_t K_,
ck_tile::index_t stride_A_,
ck_tile::index_t stride_B_,
ck_tile::index_t stride_C_)
: GemmHostArgs(a_ptr_, b_ptr_, c_ptr_, 1, M_, N_, K_, stride_A_, stride_B_, stride_C_),
NumTokens(NumTokens_),
TopK(TopK_),
p_sorted_token_ids(p_sorted_token_ids_),
p_sorted_expert_ids(p_sorted_expert_ids_),
p_max_token_id(p_max_token_id_)
{
}
private:
static constexpr index_t KBatch = 1;
};
template <typename TilePartitioner_,
typename GemmPipeline_,
typename EpiloguePipeline_,
bool IsInputGemm_ = true>
struct MoeGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, EpiloguePipeline_>
{
using TilePartitioner = remove_cvref_t<TilePartitioner_>;
using GemmPipeline = remove_cvref_t<GemmPipeline_>;
using EpiloguePipeline = remove_cvref_t<EpiloguePipeline_>;
using ALayout = remove_cvref_t<typename GemmPipeline::ALayout>;
using BLayout = remove_cvref_t<typename GemmPipeline::BLayout>;
using CLayout = remove_cvref_t<typename GemmPipeline::CLayout>;
static constexpr bool IsInputGemm = IsInputGemm_;
using ADataType = remove_cvref_t<typename GemmPipeline::ADataType>;
using BDataType = remove_cvref_t<typename GemmPipeline::BDataType>;
using CDataType = remove_cvref_t<typename EpiloguePipeline::ODataType>;
using OffsetTile1DPartitioner = OffsettedTile1DPartitioner<TilePartitioner>;
using Base = GemmKernel<TilePartitioner_, GemmPipeline_, EpiloguePipeline_>;
using GemmKernelArgs = typename Base::GemmKernelArgs;
using SplitKBatchOffset = typename Base::SplitKBatchOffset;
static constexpr index_t KernelBlockSize = GemmPipeline::BlockSize;
static constexpr auto I0 = number<0>();
static constexpr auto I1 = number<1>();
static constexpr auto I2 = number<2>();
struct MoeGemmKernelArgs : public GemmKernelArgs
{
const ck_tile::index_t* p_sorted_token_ids;
const ck_tile::index_t* p_sorted_expert_ids;
const ck_tile::index_t* p_max_token_id;
ck_tile::index_t NumTokens;
ck_tile::index_t TopK;
CK_TILE_HOST MoeGemmKernelArgs() noexcept = default;
CK_TILE_HOST MoeGemmKernelArgs(const ck_tile::index_t* p_sorted_token_ids_,
const ck_tile::index_t* p_sorted_expert_ids_,
const ck_tile::index_t* p_max_token_id_,
const void* a_ptr_,
const void* b_ptr_,
void* c_ptr_,
ck_tile::index_t NumTokens_,
ck_tile::index_t TopK_,
ck_tile::index_t M_,
ck_tile::index_t N_,
ck_tile::index_t K_,
ck_tile::index_t stride_A_,
ck_tile::index_t stride_B_,
ck_tile::index_t stride_C_,
ck_tile::index_t KBatch)
: GemmKernelArgs{a_ptr_,
b_ptr_,
c_ptr_,
M_,
N_,
K_,
stride_A_,
stride_B_,
stride_C_,
KBatch},
p_sorted_token_ids(p_sorted_token_ids_),
p_sorted_expert_ids(p_sorted_expert_ids_),
p_max_token_id(p_max_token_id_),
NumTokens(NumTokens_),
TopK(TopK_)
{
}
CK_TILE_HOST static constexpr MoeGemmKernelArgs
MakeKernelArgs(const MoeGemmHostArgs& hostArgs)
{
printf("in moe gemm kernel args!");
return MoeGemmKernelArgs{hostArgs.p_sorted_token_ids,
hostArgs.p_sorted_expert_ids,
hostArgs.p_max_token_id,
hostArgs.a_ptr,
hostArgs.b_ptr,
hostArgs.c_ptr,
hostArgs.NumTokens,
hostArgs.TopK,
hostArgs.M,
hostArgs.N,
hostArgs.K,
hostArgs.stride_A,
hostArgs.stride_B,
hostArgs.stride_C,
1
/*hostArgs.k_batch*/};
}
};
[[nodiscard]] CK_TILE_HOST static const std::string GetName()
{
// clang-format off
using P_ = GemmPipeline;
return concat('_', "moe_gemm", gemm_prec_str<ADataType, BDataType>,
concat('x', P_::MPerBlock, P_::NPerBlock, P_::KPerBlock),
concat('x', P_::GetVectorSizeA(), P_::GetVectorSizeB(), P_::GetVectorSizeC()),
concat('x', P_::kPadM, P_::kPadN, P_::kPadK));
// clang-format on
}
__host__ static constexpr auto BlockSize() -> dim3 { return dim3(KernelBlockSize); }
__host__ static constexpr auto GridSize(index_t M, index_t N, index_t KBatch)
{
// TODO: remove assertion
assert(KBatch == 1);
return Base::GridSize(M, N, KBatch);
}
CK_TILE_HOST_DEVICE static constexpr auto GetSmemSize() -> index_t
{
return max(GemmPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize());
}
template <memory_operation_enum DstInMemOp = memory_operation_enum::set>
CK_TILE_DEVICE static auto MakeGemmTensorViews(const ADataType* a_ptr,
const BDataType* b_ptr,
CDataType* c_ptr,
const MoeGemmKernelArgs& kargs,
const SplitKBatchOffset& splitk_batch_offset)
{
static_assert(!TilePartitioner::BlockGemmShape::PermuteA, "Not implemented!");
const auto& a_tensor_view = [&]() {
if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
{
return make_naive_tensor_view<address_space_enum::global>(
a_ptr,
make_tuple(IsInputGemm ? kargs.NumTokens : kargs.NumTokens * kargs.TopK,
splitk_batch_offset.splitted_k),
make_tuple(kargs.stride_A, 1),
number<GemmPipeline::GetVectorSizeA()>{},
number<1>{});
}
else
{
return make_naive_tensor_view<address_space_enum::global>(
a_ptr,
make_tuple(splitk_batch_offset.splitted_k,
IsInputGemm ? kargs.NumTokens : kargs.NumTokens * kargs.TopK),
make_tuple(kargs.stride_A, 1),
number<GemmPipeline::GetVectorSizeA()>{},
number<1>{});
}
}();
const auto& b_tensor_view = [&]() {
if constexpr(std::is_same_v<BLayout, tensor_layout::gemm::RowMajor>)
{
if constexpr(TilePartitioner::BlockGemmShape::PermuteB)
{
constexpr index_t K1 = GemmPipeline::GetSmemPackB();
const index_t K0 = splitk_batch_offset.splitted_k / K1;
constexpr index_t VectorSizeB = std::min(K1, GemmPipeline::GetVectorSizeB());
const auto b_k0_n_k1_desc =
make_naive_tensor_descriptor(make_tuple(K0, kargs.N, K1),
make_tuple(kargs.N * K1, K1, I1),
number<VectorSizeB>{},
number<1>{});
const auto b_n_k_desc = transform_tensor_descriptor(
b_k0_n_k1_desc,
make_tuple(make_merge_transform(make_tuple(K0, K1)),
make_pass_through_transform(kargs.N)),
make_tuple(sequence<0, 2>{}, sequence<1>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
return make_tensor_view<address_space_enum::global>(b_ptr, b_n_k_desc);
}
else
{
return make_naive_tensor_view<address_space_enum::global>(
b_ptr,
make_tuple(splitk_batch_offset.splitted_k, kargs.N),
make_tuple(kargs.stride_B, 1),
number<GemmPipeline::GetVectorSizeB()>{},
number<1>{});
}
}
else
{
if constexpr(TilePartitioner::BlockGemmShape::PermuteB)
{
constexpr index_t K1 = GemmPipeline::GetSmemPackB();
const index_t K0 = splitk_batch_offset.splitted_k / K1;
constexpr index_t VectorSizeB = std::min(K1, GemmPipeline::GetVectorSizeB());
const auto b_k0_n_k1_desc =
make_naive_tensor_descriptor(make_tuple(K0, kargs.N, K1),
make_tuple(kargs.N * K1, K1, I1),
number<VectorSizeB>{},
number<1>{});
const auto b_n_k_desc = transform_tensor_descriptor(
b_k0_n_k1_desc,
make_tuple(make_merge_transform(make_tuple(K0, K1)),
make_pass_through_transform(kargs.N)),
make_tuple(sequence<0, 2>{}, sequence<1>{}),
make_tuple(sequence<1>{}, sequence<0>{}));
return make_tensor_view<address_space_enum::global>(b_ptr, b_n_k_desc);
}
else
{
return make_naive_tensor_view<address_space_enum::global>(
b_ptr,
make_tuple(kargs.N, splitk_batch_offset.splitted_k),
make_tuple(kargs.stride_B, 1),
number<GemmPipeline::GetVectorSizeB()>{},
number<1>{});
}
}
}();
// TODO: enable vector write for C in ColMajor
const auto& c_tensor_view = [&]() {
if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>)
{
return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
c_ptr,
make_tuple(IsInputGemm ? kargs.NumTokens * kargs.TopK : kargs.NumTokens,
kargs.N),
make_tuple(kargs.stride_C, 1),
number<EpiloguePipeline::GetVectorSizeC()>{},
number<1>{});
}
else
{
return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
c_ptr,
make_tuple(IsInputGemm ? kargs.NumTokens * kargs.TopK : kargs.NumToken,
kargs.N),
make_tuple(1, kargs.stride_C),
number<1>{},
number<1>{});
}
}();
return make_tuple(a_tensor_view, b_tensor_view, c_tensor_view);
}
template <typename AView>
CK_TILE_DEVICE static auto GetATransformGemmView(const AView& view, const index_t token_id)
{
if constexpr(std::is_same_v<tensor_layout::gemm::RowMajor, ALayout>)
return transform_tensor_view(
view,
make_tuple(make_indexing_transform(
view.get_tensor_descriptor().get_length(number<0>()), token_id),
make_pass_through_transform(
view.get_tensor_descriptor().get_length(number<1>()))),
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
else
return transform_tensor_view(
view,
make_tuple(make_pass_through_transform(
view.get_tensor_descriptor().get_length(number<0>())),
make_indexing_transform(
view.get_tensor_descriptor().get_length(number<1>()), token_id)),
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
}
template <typename CView>
CK_TILE_DEVICE static auto GetCTransformGemmView(const CView& view, const index_t token_id)
{
if constexpr(std::is_same_v<tensor_layout::gemm::RowMajor, CLayout>)
return transform_tensor_view(
view,
make_tuple(make_indexing_transform(
view.get_tensor_descriptor().get_length(number<0>()), token_id),
make_pass_through_transform(
view.get_tensor_descriptor().get_length(number<1>()))),
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
else
return transform_tensor_view(
view,
make_tuple(make_pass_through_transform(
view.get_tensor_descriptor().get_length(number<0>())),
make_indexing_transform(
view.get_tensor_descriptor().get_length(number<1>()), token_id)),
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
}
template <typename PadView>
CK_TILE_DEVICE static auto TransformGemmPadViews(const PadView& views, const index_t token_id)
{
auto a_pad_view = views.at(number<0>());
auto b_pad_view = views.at(number<1>());
auto c_pad_view = views.at(number<2>());
const auto a_gather_view = GetATransformGemmView(a_pad_view, token_id);
// TODO： Caculate expert offset of the buf in B.
// const auto c_scatter_view = GetCTransformGemmView(c_pad_view, token_id);
// if (token_id){}
return make_tuple(a_gather_view, b_pad_view, c_pad_view);
}
template <typename PadView>
CK_TILE_DEVICE static auto
MakeGemmTileWindows(const PadView& views, const index_t i_m, const index_t i_n)
{
const auto& a_pad_view = views.at(number<0>{});
const auto& b_pad_view = views.at(number<1>{});
const auto& c_pad_view = views.at(number<2>{});
if(i_m) {}
const auto& a_block_window = [&]() {
if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
{
return make_tile_window(a_pad_view,
make_tuple(number<TilePartitioner::MPerBlock>{},
number<TilePartitioner::KPerBlock>{}),
{0, 0});
}
else
{
return make_tile_window(a_pad_view,
make_tuple(number<TilePartitioner::KPerBlock>{},
number<TilePartitioner::MPerBlock>{}),
{0, 0});
}
}();
const auto& b_block_window = [&]() {
if constexpr(std::is_same_v<BLayout, tensor_layout::gemm::ColumnMajor>)
{
return make_tile_window(b_pad_view,
make_tuple(number<TilePartitioner::NPerBlock>{},
number<TilePartitioner::KPerBlock>{}),
{i_n, 0});
}
else
{
return make_tile_window(b_pad_view,
make_tuple(number<TilePartitioner::KPerBlock>{},
number<TilePartitioner::NPerBlock>{}),
{0, i_n});
}
}();
auto c_block_window = make_tile_window(
c_pad_view,
make_tuple(number<TilePartitioner::MPerBlock>{}, number<TilePartitioner::NPerBlock>{}),
{0, i_n});
return make_tuple(a_block_window, b_block_window, c_block_window);
}
template <bool IsInputGemm = true>
CK_TILE_DEVICE void operator()(const MoeGemmKernelArgs gemm_desc) const
{
// TODO: implement C scatter store accordring to expert_id
// TODO: the branch without swizzle
const index_t max_token_id = __builtin_amdgcn_readfirstlane(gemm_desc.p_max_token_id[0]);
const index_t block_id = ck_tile::get_block_1d_id();
// TODO: check the block id caculation
const auto [expert_blk_id, _] =
OffsetTile1DPartitioner::GetOffsetedTileIndex(0, gemm_desc.M, gemm_desc.N);
if(expert_blk_id * TilePartitioner::MPerBlock >= max_token_id)
return;
const index_t NBlocks = gemm_desc.N / TilePartitioner::NPerBlock;
const index_t expert_id = gemm_desc.p_sorted_expert_ids[expert_blk_id];
const index_t prefix_blk_m = gemm_desc.p_max_token_id[1 + expert_id];
const index_t blk_cnt_of_eid = gemm_desc.p_max_token_id[2 + expert_id];
// printf("expert_blk_id: %d, expert_id: %d \n",expert_blk_id, expert_id);
// expert_id = expert_blk_id;
const index_t block_start = prefix_blk_m * NBlocks;
const index_t ecnt = blk_cnt_of_eid - prefix_blk_m;
const index_t expert_swizzle = ecnt > 0 ? ecnt : 1;
// index_t block_end = block_start + blk_cnt_of_eid * NBlocks;
const index_t block_id_start_in_expert = block_id - block_start;
const index_t im = __builtin_amdgcn_readfirstlane(prefix_blk_m + block_id_start_in_expert /
8 % expert_swizzle);
const index_t in = __builtin_amdgcn_readfirstlane(
block_id_start_in_expert % 8 + block_id_start_in_expert / (8 * expert_swizzle) * 8);
const auto a_coord = GemmPipeline::GetACoord(); // 2d thread offset, [i_row, i_col]
const auto sorted_token_id = a_coord[number<0>{}] + im * TilePartitioner::MPerBlock;
// constexpr auto AMRepeat = GemmPipeline::GetAMRepeat();
// ck_tile::statically_indexed_array<ck_tile::index_t, AMRepeat> gather_offset;
// static_for<0, AMRepeat, 1>{}([&](auto thr_offset_m){
// const index_t fused_token = gemm_desc.p_sorted_token_ids[sorted_token_id +
// thr_offset_m]; gather_offset(thr_offset_m) = fused_token & 0xffffff;
// });
const index_t fused_token = gemm_desc.p_sorted_token_ids[sorted_token_id];
// TODO: token_id should include topk offset depends on ffn1 or ffn2
const index_t token_id = fused_token & 0xffffff;
if constexpr(!IsInputGemm)
{
token_id = token_id * gemm_desc.TopK + (fused_token >> 24);
}
const index_t expert_stride = __builtin_amdgcn_readfirstlane(gemm_desc.N * gemm_desc.K);
const typename Base::SplitKBatchOffset splitk_batch_offset(gemm_desc);
// options
const ADataType* a_ptr =
static_cast<const ADataType*>(gemm_desc.a_ptr) + splitk_batch_offset.a_k_split_offset;
const BDataType* b_ptr = static_cast<const BDataType*>(gemm_desc.b_ptr) +
splitk_batch_offset.b_k_split_offset + expert_stride * expert_id;
CDataType* c_ptr = static_cast<CDataType*>(gemm_desc.c_ptr);
const auto& gemm_tensor_views_tuple =
MakeGemmTensorViews(a_ptr, b_ptr, c_ptr, gemm_desc, splitk_batch_offset);
const auto& gemm_pad_views = Base::MakeGemmPadViews(gemm_tensor_views_tuple);
const auto& transformed_views = TransformGemmPadViews(gemm_pad_views, token_id);
auto gemm_tile_windows = MakeGemmTileWindows(
transformed_views, im * TilePartitioner::MPerBlock, in * TilePartitioner::NPerBlock);
const index_t num_loop =
__builtin_amdgcn_readfirstlane(TilePartitioner::GetLoopNum(gemm_desc.K));
// printf("num_loop: %d", num_loop);
static_assert(GemmPipeline::DoubleSmemBuffer == true,
"For now, only support doublesmembuffer");
__shared__ char smem_ptr_0[GetSmemSize()];
__shared__ char smem_ptr_1[GetSmemSize()];
// Run GEMM cooperatively by whole workgroup.
const auto& a_block_window = gemm_tile_windows.at(number<0>{});
const auto& b_block_window = gemm_tile_windows.at(number<1>{});
const auto& c_block_tile = GemmPipeline{}.template operator()(
a_block_window, b_block_window, num_loop, smem_ptr_0, smem_ptr_1);
// Run Epilogue Pipeline
auto& c_block_window = gemm_tile_windows.at(number<2>{});
EpiloguePipeline{}.template operator()<decltype(c_block_window), decltype(c_block_tile)>(
c_block_window,
c_block_tile,
smem_ptr_0,
gemm_desc.p_sorted_token_ids,
im * TilePartitioner::MPerBlock);
}
};
} // namespace ck_tile

15
include/ck_tile/ops/moe_gemm.hpp Normal file
View File

@@ -0,0 +1,15 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck_tile/ops/flatmm/block/block_flatmm_asmem_bsmem_creg_v1.hpp"
#include "ck_tile/ops/flatmm/block/block_flatmm_asmem_bsmem_creg_v1_custom_policy.hpp"
#include "ck_tile/ops/flatmm/pipeline/tile_flatmm_shape.hpp"
#include "ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp"
#include "ck_tile/ops/moe_gemm/kernel/moe_gemm_kernel.hpp"
#include "ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm.hpp"
#include "ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm_policy.hpp"
#include "ck_tile/ops/common/generic_2d_block_shape.hpp"
#include "ck_tile/ops/common/tensor_layout.hpp"
#include "ck_tile/ops/common/utils.hpp"

649
include/ck_tile/ops/moe_gemm/kernel/moe_gemm_kernel.hpp Normal file
View File

@@ -0,0 +1,649 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck_tile/core/numeric/math.hpp"
#include "ck_tile/core/utility/literals.hpp"
#include "ck_tile/ops/flatmm/kernel/flatmm_kernel.hpp"
#include "ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp"
#include "ck_tile/host.hpp"
// #define disable_tile_gs
namespace ck_tile {
struct MoeGemmHostArgs : public ck_tile::FlatmmHostArgs
{
ck_tile::index_t NumTokens;
ck_tile::index_t TopK;
const ck_tile::index_t* p_sorted_token_ids;
const ck_tile::index_t* p_sorted_expert_ids;
const ck_tile::index_t* p_max_token_id;
// TODO: add kbatch for splitk
CK_TILE_HOST MoeGemmHostArgs() noexcept = default;
CK_TILE_HOST MoeGemmHostArgs(const ck_tile::index_t* p_sorted_token_ids_,
const ck_tile::index_t* p_sorted_expert_ids_,
const ck_tile::index_t* p_max_token_id_,
const void* a_ptr_,
const void* b_shuffle_ptr_,
void* c_ptr_,
ck_tile::index_t NumTokens_,
ck_tile::index_t TopK_,
ck_tile::index_t k_batch_,
ck_tile::index_t M_,
ck_tile::index_t N_,
ck_tile::index_t K_,
ck_tile::index_t stride_A_,
ck_tile::index_t stride_B_,
ck_tile::index_t stride_C_)
: FlatmmHostArgs(a_ptr_, b_shuffle_ptr_, c_ptr_, k_batch_, M_, N_, K_, stride_A_, stride_B_, stride_C_),
NumTokens(NumTokens_),
TopK(TopK_),
p_sorted_token_ids(p_sorted_token_ids_),
p_sorted_expert_ids(p_sorted_expert_ids_),
p_max_token_id(p_max_token_id_)
{
}
// TODO: why kBatch?
// private:
// static constexpr index_t KBatch = 1;
};
template <typename TilePartitioner_,
typename FlatmmPipeline_,
typename EpiloguePipeline_,
bool IsInputGemm_ = true>
struct MoeGemmKernel
{
using TilePartitioner = remove_cvref_t<TilePartitioner_>;
using FlatmmPipeline = remove_cvref_t<FlatmmPipeline_>;
using EpiloguePipeline = remove_cvref_t<EpiloguePipeline_>;
using ALayout = remove_cvref_t<typename FlatmmPipeline::ALayout>;
using BLayout = remove_cvref_t<typename FlatmmPipeline::BLayout>;
using CLayout = remove_cvref_t<typename FlatmmPipeline::CLayout>;
using BlockGemmShape =
remove_cvref_t<typename FlatmmPipeline::BlockGemmShape>; // TileFlatmmShape
static constexpr bool IsInputGemm = IsInputGemm_;
using ADataType = remove_cvref_t<typename FlatmmPipeline::ADataType>;
using BDataType = remove_cvref_t<typename FlatmmPipeline::BDataType>;
using CDataType = remove_cvref_t<typename EpiloguePipeline::ODataType>;
using OffsetTile1DPartitioner = OffsettedTile1DPartitioner<TilePartitioner>;
static constexpr index_t KernelBlockSize = FlatmmPipeline::BlockSize;
static constexpr auto I0 = number<0>();
static constexpr auto I1 = number<1>();
static constexpr auto I2 = number<2>();
struct MoeGemmKernelArgs
{
const ck_tile::index_t* p_sorted_token_ids;
const ck_tile::index_t* p_sorted_expert_ids;
const ck_tile::index_t* p_max_token_id;
const void* p_a_ptr;
const void* p_b_shuffle_ptr;
void* p_c_ptr;
ck_tile::index_t NumTokens;
ck_tile::index_t TopK;
ck_tile::index_t M;
ck_tile::index_t N;
ck_tile::index_t K;
ck_tile::index_t stride_A;
ck_tile::index_t stride_B;
ck_tile::index_t stride_C;
ck_tile::index_t k_batch;
//
// CK_TILE_HOST MoeGemmKernelArgs() noexcept = default;
// CK_TILE_HOST MoeGemmKernelArgs(const ck_tile::index_t* p_sorted_token_ids_,
// const ck_tile::index_t* p_sorted_expert_ids_,
// const ck_tile::index_t* p_max_token_id_,
// const void* a_ptr_,
// const void* b_shuffle_ptr_,
// void* c_ptr_,
// ck_tile::index_t NumTokens_,
// ck_tile::index_t TopK_,
// ck_tile::index_t M_,
// ck_tile::index_t N_,
// ck_tile::index_t K_,
// ck_tile::index_t stride_A_,
// ck_tile::index_t stride_B_,
// ck_tile::index_t stride_C_,
// ck_tile::index_t KBatch) :
// p_sorted_token_ids(p_sorted_token_ids_),
// p_sorted_expert_ids(p_sorted_expert_ids_),
// p_max_token_id(p_max_token_id_),
// p_a_ptr(a_ptr_),
// p_b_shuffle_ptr(b_shuffle_ptr_),
// p_c_ptr(c_ptr_),
// NumTokens(NumTokens_),
// TopK(TopK_),
// M(M_),
// N(N_),
// K(K_),
// stride_A(stride_A_),
// stride_B(stride_B_),
// stride_C(stride_C_),
// k_batch(KBatch)
// {
// }
};
CK_TILE_HOST static constexpr MoeGemmKernelArgs MakeKernelArgs(const MoeGemmHostArgs& hostArgs)
{
printf("in moe gemm kernel args! \n");
return MoeGemmKernelArgs{hostArgs.p_sorted_token_ids,
hostArgs.p_sorted_expert_ids,
hostArgs.p_max_token_id,
hostArgs.a_ptr,
hostArgs.b_shuffle_ptr,
hostArgs.c_ptr,
hostArgs.NumTokens,
hostArgs.TopK,
hostArgs.M,
hostArgs.N,
hostArgs.K,
hostArgs.stride_A,
hostArgs.stride_B,
hostArgs.stride_C,
1
/*hostArgs.k_batch*/};
}
[[nodiscard]] CK_TILE_HOST static const std::string GetName()
{
// clang-format off
using P_ = FlatmmPipeline;
return concat('_', "moe_gemm", gemm_prec_str<ADataType, BDataType>,
concat('x', P_::kMPerBlock, P_::kNPerBlock, P_::kKPerBlock),
concat('x', P_::GetVectorSizeA(), P_::GetVectorSizeB(), P_::GetVectorSizeC()),
concat('x', P_::kPadM, P_::kPadN, P_::kPadK));
// clang-format on
}
__host__ static constexpr auto BlockSize() -> dim3 { return dim3(KernelBlockSize); }
__host__ static constexpr auto GridSize(index_t M, index_t N, index_t KBatch)
{
// TODO: remove assertion
assert(KBatch == 1);
return dim3(TilePartitioner::GridSize(M, N), 1, KBatch);
}
CK_TILE_HOST_DEVICE static constexpr auto GetSmemSize() -> index_t
{
return max(FlatmmPipeline::GetSmemSize(), EpiloguePipeline::GetSmemSize());
}
struct SplitKBatchOffset
{
__device__ SplitKBatchOffset(const MoeGemmKernelArgs& kargs,
const std::size_t k_id = blockIdx.z)
{
constexpr auto K1 = TilePartitioner::BlockGemmShape::WarpTile::at(number<2>{});
const index_t K_t = __builtin_amdgcn_readfirstlane(kargs.k_batch * K1);
const index_t KRead = __builtin_amdgcn_readfirstlane((kargs.K + K_t - 1) / K_t * K1);
if constexpr(std::is_same_v<tensor_layout::gemm::RowMajor, ALayout>)
{
a_k_split_offset = __builtin_amdgcn_readfirstlane(k_id * KRead);
}
else if constexpr(std::is_same_v<tensor_layout::gemm::ColumnMajor, ALayout>)
{
a_k_split_offset = __builtin_amdgcn_readfirstlane(k_id * KRead * kargs.stride_A);
}
if constexpr(std::is_same_v<tensor_layout::gemm::RowMajor, BLayout>)
{
b_k_split_offset = __builtin_amdgcn_readfirstlane(k_id * KRead * kargs.stride_B);
}
else if constexpr(std::is_same_v<tensor_layout::gemm::ColumnMajor, BLayout>)
{
b_k_split_offset = __builtin_amdgcn_readfirstlane(k_id * KRead);
}
if(k_id < static_cast<uint32_t>(kargs.k_batch - 1))
{
splitted_k = __builtin_amdgcn_readfirstlane(KRead);
}
else
{
splitted_k = __builtin_amdgcn_readfirstlane(kargs.K - KRead * (kargs.k_batch - 1));
}
}
index_t a_k_split_offset;
index_t b_k_split_offset;
index_t splitted_k;
};
template <memory_operation_enum DstInMemOp = memory_operation_enum::set>
CK_TILE_DEVICE static auto MakeGemmTensorViews(const ADataType* a_ptr,
const BDataType* b_flat_ptr,
CDataType* c_ptr,
const MoeGemmKernelArgs& kargs,
const SplitKBatchOffset& splitk_batch_offset)
{
// static_assert(!TilePartitioner::BlockGemmShape::PermuteA, "Not implemented!");
const auto& a_tensor_view = [&]() {
if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
{
return make_naive_tensor_view<address_space_enum::global>(
a_ptr,
make_tuple(IsInputGemm ? kargs.NumTokens : kargs.NumTokens * kargs.TopK,
splitk_batch_offset.splitted_k),
make_tuple(kargs.stride_A, 1),
number<FlatmmPipeline::GetVectorSizeA()>{},
number<1>{});
}
else
{
return make_naive_tensor_view<address_space_enum::global>(
a_ptr,
make_tuple(splitk_batch_offset.splitted_k,
IsInputGemm ? kargs.NumTokens : kargs.NumTokens * kargs.TopK),
make_tuple(kargs.stride_A, 1),
number<FlatmmPipeline::GetVectorSizeA()>{},
number<1>{});
}
}();
index_t kFlatK = FlatmmPipeline::flatKPerWarp * (splitk_batch_offset.splitted_k /
BlockGemmShape::WarpTile::at(number<2>{}));
index_t kFlatN = kargs.N * kargs.K / kFlatK;
const auto& b_flat_tensor_view = [&]() {
return make_naive_tensor_view<address_space_enum::global>(
b_flat_ptr,
make_tuple(kFlatN, kFlatK),
make_tuple(kFlatK, 1),
number<FlatmmPipeline::GetVectorSizeB()>{},
number<1>{});
}();
// const auto& b_tensor_view = [&]() {
// if constexpr(std::is_same_v<BLayout, tensor_layout::gemm::RowMajor>)
// {
// if constexpr(TilePartitioner::BlockGemmShape::PermuteB)
// {
// constexpr index_t K1 = FlatmmPipeline::GetSmemPackB();
// const index_t K0 = splitk_batch_offset.splitted_k / K1;
// constexpr index_t VectorSizeB = std::min(K1, FlatmmPipeline::GetVectorSizeB());
// const auto b_k0_n_k1_desc =
// make_naive_tensor_descriptor(make_tuple(K0, kargs.N, K1),
// make_tuple(kargs.N * K1, K1, I1),
// number<VectorSizeB>{},
// number<1>{});
// const auto b_n_k_desc = transform_tensor_descriptor(
// b_k0_n_k1_desc,
// make_tuple(make_merge_transform(make_tuple(K0, K1)),
// make_pass_through_transform(kargs.N)),
// make_tuple(sequence<0, 2>{}, sequence<1>{}),
// make_tuple(sequence<0>{}, sequence<1>{}));
// return make_tensor_view<address_space_enum::global>(b_ptr, b_n_k_desc);
// }
// else
// {
// return make_naive_tensor_view<address_space_enum::global>(
// b_ptr,
// make_tuple(splitk_batch_offset.splitted_k, kargs.N),
// make_tuple(kargs.stride_B, 1),
// number<FlatmmPipeline::GetVectorSizeB()>{},
// number<1>{});
// }
// }
// else
// {
// if constexpr(TilePartitioner::BlockGemmShape::PermuteB)
// {
// constexpr index_t K1 = FlatmmPipeline::GetSmemPackB();
// const index_t K0 = splitk_batch_offset.splitted_k / K1;
// constexpr index_t VectorSizeB = std::min(K1, FlatmmPipeline::GetVectorSizeB());
// const auto b_k0_n_k1_desc =
// make_naive_tensor_descriptor(make_tuple(K0, kargs.N, K1),
// make_tuple(kargs.N * K1, K1, I1),
// number<VectorSizeB>{},
// number<1>{});
// const auto b_n_k_desc = transform_tensor_descriptor(
// b_k0_n_k1_desc,
// make_tuple(make_merge_transform(make_tuple(K0, K1)),
// make_pass_through_transform(kargs.N)),
// make_tuple(sequence<0, 2>{}, sequence<1>{}),
// make_tuple(sequence<1>{}, sequence<0>{}));
// return make_tensor_view<address_space_enum::global>(b_ptr, b_n_k_desc);
// }
// else
// {
// return make_naive_tensor_view<address_space_enum::global>(
// b_ptr,
// make_tuple(kargs.N, splitk_batch_offset.splitted_k),
// make_tuple(kargs.stride_B, 1),
// number<FlatmmPipeline::GetVectorSizeB()>{},
// number<1>{});
// }
// }
// }();
// TODO: enable vector write for C in ColMajor
const auto& c_tensor_view = [&]() {
if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>)
{
return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
c_ptr,
make_tuple(IsInputGemm ? kargs.NumTokens * kargs.TopK : kargs.NumTokens,
kargs.N),
make_tuple(kargs.stride_C, 1),
number<EpiloguePipeline::GetVectorSizeC()>{},
number<1>{});
}
else
{
return make_naive_tensor_view<address_space_enum::global, DstInMemOp>(
c_ptr,
make_tuple(IsInputGemm ? kargs.NumTokens * kargs.TopK : kargs.NumToken,
kargs.N),
make_tuple(1, kargs.stride_C),
number<1>{},
number<1>{});
}
}();
return make_tuple(a_tensor_view, b_flat_tensor_view, c_tensor_view);
}
template <typename TensorView>
CK_TILE_DEVICE static auto MakeGemmPadViews(const TensorView& views)
{
const auto& a_pad_view = [&]() {
const auto& a_tensor_view = views.at(I0);
if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
{
return pad_tensor_view(a_tensor_view,
make_tuple(number<TilePartitioner::MPerBlock>{},
number<TilePartitioner::KPerBlock>{}),
sequence<false, FlatmmPipeline::kPadK>{});
}
else
{
return pad_tensor_view(a_tensor_view,
make_tuple(number<TilePartitioner::KPerBlock>{},
number<TilePartitioner::MPerBlock>{}),
sequence<false, FlatmmPipeline::kPadM>{});
}
}();
const auto& b_flat_tensor_view = views.at(I1);
// TODO vector write in for C in ColMajor
const auto& c_pad_view = [&]() {
const auto& c_tensor_view = views.at(I2);
if constexpr(std::is_same_v<CLayout, tensor_layout::gemm::RowMajor>)
{
return pad_tensor_view(c_tensor_view,
make_tuple(number<TilePartitioner::MPerBlock>{},
number<TilePartitioner::NPerBlock>{}),
sequence<false, FlatmmPipeline::kPadN>{});
}
else
{
return pad_tensor_view(c_tensor_view,
make_tuple(number<TilePartitioner::MPerBlock>{},
number<TilePartitioner::NPerBlock>{}),
sequence<FlatmmPipeline::kPadM, false>{});
}
}();
return make_tuple(a_pad_view, b_flat_tensor_view, c_pad_view);
}
template <typename AView>
CK_TILE_DEVICE static auto GetATransformGemmView(const AView& view, const index_t token_id)
{
if constexpr(std::is_same_v<tensor_layout::gemm::RowMajor, ALayout>)
return transform_tensor_view(
view,
make_tuple(make_indexing_transform(
view.get_tensor_descriptor().get_length(number<0>()), token_id),
make_pass_through_transform(
view.get_tensor_descriptor().get_length(number<1>()))),
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
else
return transform_tensor_view(
view,
make_tuple(make_pass_through_transform(
view.get_tensor_descriptor().get_length(number<0>())),
make_indexing_transform(
view.get_tensor_descriptor().get_length(number<1>()), token_id)),
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
}
// template <typename CView>
// CK_TILE_DEVICE static auto GetCTransformGemmView(const CView& view, const index_t token_id)
// {
// if constexpr(std::is_same_v<tensor_layout::gemm::RowMajor, CLayout>)
// return transform_tensor_view(
// view,
// make_tuple(make_indexing_transform(
// view.get_tensor_descriptor().get_length(number<0>()), token_id),
// make_pass_through_transform(
// view.get_tensor_descriptor().get_length(number<1>()))),
// make_tuple(sequence<0>{}, sequence<1>{}),
// make_tuple(sequence<0>{}, sequence<1>{}));
// else
// return transform_tensor_view(
// view,
// make_tuple(make_pass_through_transform(
// view.get_tensor_descriptor().get_length(number<0>())),
// make_indexing_transform(
// view.get_tensor_descriptor().get_length(number<1>()), token_id)),
// make_tuple(sequence<0>{}, sequence<1>{}),
// make_tuple(sequence<0>{}, sequence<1>{}));
// }
template <typename PadView>
CK_TILE_DEVICE static auto TransformGemmPadViews(const PadView& views, const index_t token_id)
{
auto a_pad_view = views.at(number<0>());
auto b_pad_view = views.at(number<1>());
auto c_pad_view = views.at(number<2>());
const auto a_gather_view = GetATransformGemmView(a_pad_view, token_id);
// TODO： Caculate expert offset of the buf in B.
// const auto c_scatter_view = GetCTransformGemmView(c_pad_view, token_id);
// if (token_id){}
return make_tuple(a_gather_view, b_pad_view, c_pad_view);
}
template <typename PadView>
CK_TILE_DEVICE static auto
MakeGemmTileWindows(const PadView& views, const index_t i_m, const index_t i_n)
{
(void)i_m;
const auto& a_pad_view = views.at(number<0>{});
const auto& b_flat_pad_view = views.at(number<1>{});
const auto& c_pad_view = views.at(number<2>{});
// if(i_m) {}
const auto& a_block_window = [&]() {
if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::RowMajor>)
{
return make_tile_window(a_pad_view,
make_tuple(number<TilePartitioner::MPerBlock>{},
number<TilePartitioner::KPerBlock>{}),
{0, 0});
}
else
{
return make_tile_window(a_pad_view,
make_tuple(number<TilePartitioner::KPerBlock>{},
number<TilePartitioner::MPerBlock>{}),
{0, 0});
}
}();
const auto& b_flat_block_window =
make_tile_window(b_flat_pad_view,
make_tuple(number<FlatmmPipeline::flatNPerWarp>{},
number<FlatmmPipeline::flatKPerWarp>{}),
{static_cast<int>(i_n / BlockGemmShape::WarpTile::at(I1)), 0});
auto c_block_window = make_tile_window(
c_pad_view,
make_tuple(number<TilePartitioner::MPerBlock>{}, number<TilePartitioner::NPerBlock>{}),
// {i_m, i_n});
{0, i_n});
return make_tuple(a_block_window, b_flat_block_window, c_block_window);
}
template <bool IsInputGemm = true>
CK_TILE_DEVICE void operator()(const MoeGemmKernelArgs gemm_desc) const
{
// TODO: implement C scatter store accordring to expert_id
// TODO: the branch without swizzle
const index_t max_token_id = __builtin_amdgcn_readfirstlane(gemm_desc.p_max_token_id[0]);
const index_t block_id = ck_tile::get_block_1d_id();
// TODO: check the block id caculation
const auto [expert_blk_id, _] =
OffsetTile1DPartitioner::GetOffsetedTileIndex(0, gemm_desc.M, gemm_desc.N);
if(expert_blk_id * TilePartitioner::MPerBlock >= max_token_id)
return;
const index_t NBlocks = gemm_desc.N / TilePartitioner::NPerBlock;
const index_t expert_id = gemm_desc.p_sorted_expert_ids[expert_blk_id];
const index_t prefix_blk_m = gemm_desc.p_max_token_id[1 + expert_id];
const index_t blk_cnt_of_eid = gemm_desc.p_max_token_id[2 + expert_id];
// printf("expert_blk_id: %d, expert_id: %d \n",expert_blk_id, expert_id);
// expert_id = expert_blk_id;
const index_t block_start = prefix_blk_m * NBlocks;
const index_t ecnt = blk_cnt_of_eid - prefix_blk_m;
const index_t expert_swizzle = ecnt > 0 ? ecnt : 1;
// index_t block_end = block_start + blk_cnt_of_eid * NBlocks;
const index_t block_id_start_in_expert = block_id - block_start;
const index_t im = __builtin_amdgcn_readfirstlane(prefix_blk_m + block_id_start_in_expert /
8 % expert_swizzle);
const index_t in = __builtin_amdgcn_readfirstlane(
block_id_start_in_expert % 8 + block_id_start_in_expert / (8 * expert_swizzle) * 8);
const auto a_coord = FlatmmPipeline::GetACoord(); // 2d thread offset, [i_row, i_col]
#ifdef disable_tile_gs
const auto sorted_token_id = a_coord[number<0>{}] + im * TilePartitioner::MPerBlock;
const index_t fused_token = gemm_desc.p_sorted_token_ids[sorted_token_id];
// TODO: token_id should include topk offset depends on ffn1 or ffn2
constexpr index_t token_id_mask = 0xffffff;
index_t token_id = fused_token & token_id_mask;
if constexpr(!IsInputGemm)
{
constexpr index_t token_id_offset = 24;
token_id = token_id * gemm_desc.TopK + (fused_token >> token_id_offset);
}
#else
constexpr ck_tile::index_t MRepeat = FlatmmPipeline::GetAMRepeat();
statically_indexed_array<ck_tile::index_t, MRepeat> a_offsets;
constexpr index_t token_id_mask = 0xffffff;
constexpr index_t token_id_offset = 24;
// constexpr auto kMWave = TilePartitioner::BlockGemmShape::BlockWarps::at(I0);
// constexpr auto kNWave = TilePartitioner::BlockGemmShape::BlockWarps::at(I1);
// const index_t iMWarp = get_warp_id() / kNWave;
static_for<0, MRepeat, 1>{}([&](auto m0) {
// const auto sorted_token_id = a_coord[I0] + im * TilePartitioner::MPerBlock +
// iMWarp * TilePartitioner::MPerBlock / kMWave +
// m0 * TilePartitioner::MPerBlock / kMWave / MRepeat;
const auto sorted_token_id = a_coord[I0] + im * TilePartitioner::MPerBlock +
m0 * TilePartitioner::MPerBlock / MRepeat;
const index_t fused_token = gemm_desc.p_sorted_token_ids[sorted_token_id];
// TODO: token_id should include topk offset depends on ffn1 or ffn2
index_t gather_token_id = fused_token & token_id_mask;
if constexpr(!IsInputGemm)
{
gather_token_id = gather_token_id * gemm_desc.TopK + (fused_token >> token_id_offset);
}
a_offsets[m0] = gather_token_id * gemm_desc.stride_A;
});
#endif
const index_t expert_stride = __builtin_amdgcn_readfirstlane(gemm_desc.N * gemm_desc.K);
const SplitKBatchOffset splitk_batch_offset(gemm_desc);
// options
const ADataType* a_ptr =
static_cast<const ADataType*>(gemm_desc.p_a_ptr) + splitk_batch_offset.a_k_split_offset;
const BDataType* b_shuffle_ptr = static_cast<const BDataType*>(gemm_desc.p_b_shuffle_ptr) +
splitk_batch_offset.b_k_split_offset + expert_stride * expert_id;
CDataType* c_ptr = static_cast<CDataType*>(gemm_desc.p_c_ptr);
const auto& gemm_tensor_views_tuple =
MakeGemmTensorViews(a_ptr, b_shuffle_ptr, c_ptr, gemm_desc, splitk_batch_offset);
const auto& gemm_pad_views = MakeGemmPadViews(gemm_tensor_views_tuple);
#ifdef disable_tile_gs
const auto& transformed_views = TransformGemmPadViews(gemm_pad_views, token_id);
auto gemm_tile_windows = MakeGemmTileWindows(
transformed_views, im * TilePartitioner::MPerBlock, in * TilePartitioner::NPerBlock);
#else
auto gemm_tile_windows = MakeGemmTileWindows(
gemm_pad_views, im * TilePartitioner::MPerBlock, in * TilePartitioner::NPerBlock);
#endif
const index_t num_loop =
__builtin_amdgcn_readfirstlane(TilePartitioner::GetLoopNum(gemm_desc.K));
// printf("num_loop: %d", num_loop);
// static_assert(FlatmmPipeline::DoubleSmemBuffer == true,
// "For now, only support doublesmembuffer");
__shared__ char smem_ptr_0[GetSmemSize()];
// __shared__ char smem_ptr_1[GetSmemSize()];
// Run GEMM cooperatively by whole workgroup.
const auto& a_block_window = gemm_tile_windows.at(number<0>{});
const auto& b_block_window = gemm_tile_windows.at(number<1>{});
#ifdef disable_tile_gs
const auto& c_block_tile = FlatmmPipeline{}.template operator()(
a_block_window, b_block_window, num_loop, smem_ptr_0);
#else
auto a_gather_block_tile = ck_tile::make_tile_scatter_gather(
a_block_window.get_bottom_tensor_view(),
a_block_window.get_window_lengths(),
a_block_window.get_window_origin(),
FlatmmPipeline::GetADramTileDistribution(),
a_offsets); // K DRAM tile window for
const auto& c_block_tile = FlatmmPipeline{}.template operator()(
a_gather_block_tile, b_block_window, num_loop, smem_ptr_0);
#endif
// Run Epilogue Pipeline
auto& c_block_window = gemm_tile_windows.at(number<2>{});
EpiloguePipeline{}.template operator()<decltype(c_block_window), decltype(c_block_tile)>(
c_block_window,
c_block_tile,
smem_ptr_0,
gemm_desc.p_sorted_token_ids,
im * TilePartitioner::MPerBlock,
gemm_desc.TopK,
gemm_desc.stride_C);
}
};
} // namespace ck_tile

0
include/ck_tile/ops/gemm/pipeline/moe_gemm_pipeline_ag_bg_cr.hpp → include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_ag_bg_cr.hpp
View File

5
include/ck_tile/ops/gemm/pipeline/moe_gemm_pipeline_ag_bg_cr_policy.hpp → include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_ag_bg_cr_policy.hpp
View File

@@ -6,10 +6,11 @@
#include "ck_tile/core.hpp"
#include "ck_tile/ops/common.hpp"
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp"
#include "ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp"
namespace ck_tile {
template <typename Problem>
struct MoeGemmPipelineAgBgCrPolicy : public GemmPipelineAgBgCrCompV4DefaultPolicy
struct MoeGemmPipelineAgBgCrPolicy : public UniversalFlatmmPipelineAgBgCrPolicy
{
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto MakeGlobalTileDistribution_C()
@@ -33,4 +34,4 @@ struct MoeGemmPipelineAgBgCrPolicy : public GemmPipelineAgBgCrCompV4DefaultPolic
return c_block_dstr;
}
}
} // namespace ck_tile
} // namespace ck_tile

245
include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm.hpp Normal file
View File

@@ -0,0 +1,245 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck_tile/core.hpp"
#include "ck_tile/ops/common.hpp"
#include "ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm_policy.hpp"
#include <cwchar>
namespace ck_tile {
template <typename Problem, typename PipelinePolicy = UniversalFlatmmPipelineAgBgCrPolicy>
struct MoeGemmPipelineAgBgCrImpl
{
using ADataType = remove_cvref_t<typename Problem::ADataType>;
using BDataType = remove_cvref_t<typename Problem::BDataType>;
using CDataType = remove_cvref_t<typename Problem::CDataType>;
using BlockGemmShape = remove_cvref_t<typename Problem::BlockGemmShape>;
static_assert(!std::is_same_v<BDataType, pk_int4_t>, "Not implemented");
static constexpr index_t APackedSize =
ck_tile::numeric_traits<remove_cvref_t<ADataType>>::PackedSize;
static constexpr index_t BPackedSize =
ck_tile::numeric_traits<remove_cvref_t<BDataType>>::PackedSize;
using ALayout = remove_cvref_t<typename Problem::ALayout>;
using BLayout = remove_cvref_t<typename Problem::BLayout>;
using CLayout = remove_cvref_t<typename Problem::CLayout>;
using BlockFlatmm = remove_cvref_t<decltype(PipelinePolicy::template GetBlockFlatmm<Problem>())>;
using I0 = number<0>;
using I1 = number<1>;
using I2 = number<2>;
static constexpr index_t BlockSize = Problem::kBlockSize;
static constexpr index_t kMPerBlock = BlockGemmShape::kM;
static constexpr index_t kNPerBlock = BlockGemmShape::kN;
static constexpr index_t kKPerBlock = BlockGemmShape::kK;
static constexpr index_t flatKPerWarp = BlockGemmShape::flatKPerWarp;
static constexpr index_t flatNPerWarp = BlockGemmShape::flatNPerWarp;
// static constexpr index_t GetVectorSizeA() { return PipelinePolicy::template GetVectorSizeA<Problem>(); }
// static constexpr index_t GetVectorSizeB() { return PipelinePolicy::template GetVectorSizeB<Problem>(); }
// static constexpr index_t GetVectorSizeC() { return PipelinePolicy::template GetVectorSizeC<Problem>(); }
static constexpr index_t GetVectorSizeA() { return Problem::VectorSizeA; }
static constexpr index_t GetVectorSizeB() { return Problem::VectorSizeB; }
static constexpr index_t GetVectorSizeC() { return Problem::VectorSizeC; }
static constexpr index_t GetSmemPackA() { return PipelinePolicy::template GetSmemPackA<Problem>(); }
static constexpr index_t GetSmemPackB() { return PipelinePolicy::template GetSmemPackB<Problem>(); }
static constexpr bool kPadM = Problem::kPadM;
static constexpr bool kPadN = Problem::kPadN;
static constexpr bool kPadK = Problem::kPadK;
// static constexpr bool DoubleSmemBuffer = Problem::DoubleSmemBuffer;
static constexpr bool HasHotLoop = Problem::HasHotLoop;
static constexpr auto Scheduler = Problem::Scheduler;
CK_TILE_HOST_DEVICE static constexpr auto TransposeC() { return Problem::TransposeC; }
CK_TILE_HOST_DEVICE static constexpr auto GetADramTileDistribution() {
return PipelinePolicy::template MakeADramTileDistribution<Problem>();
}
CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()
{
return PipelinePolicy::template GetSmemSize<Problem>();
}
CK_TILE_HOST_DEVICE constexpr static auto GetACoord()
{
constexpr auto a_dist = PipelinePolicy::template MakeADramTileDistribution<Problem>();
return a_dist.calculate_index();
}
// get thread coordinate of A in the threadblock
CK_TILE_HOST_DEVICE constexpr static auto GetAMRepeat()
{
constexpr auto a_dist = PipelinePolicy::template MakeADramTileDistribution<Problem>();
using ADstrEncode = typename decltype(a_dist)::DstrEncode;
constexpr ck_tile::index_t MRepeat = ADstrEncode::hs_lengthss_[number<0>{}][number<0>{}];
return MRepeat;
}
template <typename ADramBlockWindow, typename BFlatBlockWindowTmp, typename AElementFunction>
CK_TILE_HOST_DEVICE auto operator()(ADramBlockWindow& a_dram_block_window,
const AElementFunction& a_element_func,
const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp,
index_t num_loop,
void* p_smem) const
{
static_assert(
std::is_same_v<ADataType, remove_cvref_t<typename ADramBlockWindow::DataType>>,
"wrong!");
static_assert(kMPerBlock == ADramBlockWindow{}.get_window_lengths()[number<0>{}],
"wrong!");
static_assert(kKPerBlock == ADramBlockWindow{}.get_window_lengths()[number<1>{}],
"wrong!");
// A tile in LDS
ADataType* p_a_lds = static_cast<ADataType*>(p_smem);
constexpr auto a_lds_block_desc =
PipelinePolicy::template MakeALdsBlockDescriptor<Problem>();
auto a_lds_block = make_tensor_view<address_space_enum::lds>(p_a_lds, a_lds_block_desc);
// auto a_dist = PipelinePolicy::template MakeADramTileDistribution<Problem>();
// auto a_coord = a_dist.calculate_index();
// using ADstrEncode = typename decltype(a_dist)::DstrEncode;
// constexpr ck_tile::index_t MRepeat = ADstrEncode::hs_lengthss_[I0][I0];
// statically_indexed_array<ck_tile::index_t, NRepeat> a_offsets;
// static_for<0, MRepeat, 1>{}([&](auto n0) {
// int32_t seqlen_k_idx_per_repeat = cur_seqlen_k_idx + k_coord[0] + Traits::kBlockN / NRepeat * n0.value;
// int32_t page_idx = seqlen_k_idx_per_repeat / page_block_size;
// int32_t seq_idx = seqlen_k_idx_per_repeat % page_block_size;
// k_offsets[n0] = (block_indices[page_idx] * page_block_size + seq_idx) * stride_s_k;
// });
//
// // A DRAM tile window for load
// auto a_dram_tile = ck_tile::make_tile_scatter_gather(
// a_dram_block_window_tmp.get_bottom_tensor_view(),
// a_dram_block_window_tmp.get_window_lengths(),
// a_dram_block_window_tmp.get_window_origin(),
// a_dist,
// k_offsets); // K DRAM tile window for
// auto a_copy_dram_window =
// make_tile_window(a_dram_block_window_tmp.get_bottom_tensor_view(),
// make_tuple(number<kMPerBlock>{}, number<kKPerBlock>{}),
// a_dram_block_window_tmp.get_window_origin(),
// PipelinePolicy::template MakeADramTileDistribution<Problem>());
// A LDS tile window for store
auto a_copy_lds_window = make_tile_window(
a_lds_block, make_tuple(number<kMPerBlock>{}, number<kKPerBlock>{}), {0, 0});
// A LDS tile for block GEMM
auto a_lds_gemm_window = make_tile_window(
a_lds_block, make_tuple(number<kMPerBlock>{}, number<kKPerBlock>{}), {0, 0});
// Block GEMM
auto block_flatmm = BlockFlatmm();
// B flat DRAM window for load
auto b_flat_distribution =
PipelinePolicy::template MakeBFlatDramTileDistribution<Problem>();
auto b_flat_dram_window = // tile_window_with_static_distribution
make_tile_window(
b_flat_dram_block_window_tmp.get_bottom_tensor_view(), // from kernel gemm_pad_views
make_tuple(number<flatNPerWarp>{}, number<flatKPerWarp>{}),
b_flat_dram_block_window_tmp.get_window_origin(),
b_flat_distribution);
// Acc register tile
auto c_block_tile = decltype(block_flatmm(a_lds_gemm_window, b_flat_dram_window)){};
// prefetch
// global read 0
auto a_block_tile = a_dram_block_window.load();
{
// move to 1
move_tile_window(a_dram_block_window, {0, kKPerBlock});
// initialize C
tile_elementwise_inout([](auto& c) { c = 0; }, c_block_tile);
// LDS write 0
if constexpr(std::is_same_v<ALayout, tensor_layout::gemm::ColumnMajor>)
{
auto a_shuffle_tmp = make_static_distributed_tensor<ADataType>(
PipelinePolicy::template MakeShuffledARegBlockDistribution<Problem>());
shuffle_tile(a_shuffle_tmp, a_block_tile);
const auto a_block_tile_tmp = tile_elementwise_in(a_element_func, a_shuffle_tmp);
store_tile(a_copy_lds_window, a_block_tile_tmp);
}
else
{
store_tile(a_copy_lds_window, tile_elementwise_in(a_element_func, a_block_tile));
}
}
index_t iCounter = num_loop - 1;
while(iCounter > 0)
{
// global read i + 1
a_dram_block_window.load(a_block_tile);
block_sync_lds();
// GEMM i
block_flatmm(c_block_tile, a_lds_gemm_window, b_flat_dram_window);
block_sync_lds();
// move to i + 2
move_tile_window(a_dram_block_window, {0, kKPerBlock});
// LDS write i + 1
const auto a_block_tile_tmp = tile_elementwise_in(a_element_func, a_block_tile);
store_tile(a_copy_lds_window, a_block_tile_tmp);
// move to next flat K
move_tile_window(b_flat_dram_window, {0, BlockGemmShape::flatKPerBlock});
iCounter--;
}
// tail
{
block_sync_lds();
// GEMM num_loop - 1
block_flatmm(c_block_tile, a_lds_gemm_window, b_flat_dram_window);
}
return c_block_tile;
}
template <typename ADramBlockWindow, typename BFlatBlockWindowTmp>
CK_TILE_DEVICE auto operator()(ADramBlockWindow& a_dram_block_window_tmp,
const BFlatBlockWindowTmp& b_flat_dram_block_window_tmp,
index_t num_loop,
void* p_smem) const
{
return operator()(
a_dram_block_window_tmp,
[](const ADataType& a) { return a; },
b_flat_dram_block_window_tmp,
num_loop,
p_smem);
}
};
} // namespace ck_tile

37
include/ck_tile/ops/moe_gemm/pipeline/moe_gemm_pipeline_agmem_bgmem_creg_flatmm_policy.hpp Normal file
View File

@@ -0,0 +1,37 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck_tile/core.hpp"
#include "ck_tile/ops/common.hpp"
#include "ck_tile/ops/flatmm/pipeline/flatmm_pipeline_agmem_bgmem_creg_v1_policy.hpp"
namespace ck_tile {
struct MoeGemmPipelineAgBgCrPolicy : public UniversalFlatmmPipelineAgBgCrPolicy
{
template <typename Problem>
CK_TILE_HOST_DEVICE static constexpr auto MakeGlobalTileDistribution_C()
{
using S_ = remove_cvref_t<typename Problem::BlockShape>;
using WarpGemm = remove_cvref_t<typename Problem::W>;
// using CDataType = typename WarpGemm::CDataType;
constexpr auto c_block_outer_dstr_encoding =
tile_distribution_encoding<sequence<>,
tuple<sequence<S_::Repeat_M1, S_::WarpPerBlock_M1>,
sequence<S_::Repeat_N1, S_::WarpPerBlock_N1>>,
tuple<sequence<1, 2>>,
tuple<sequence<1, 1>>,
sequence<1, 2>,
sequence<0, 0>>{};
constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
c_block_outer_dstr_encoding, typename WarpGemm::CWarpDstrEncoding{});
constexpr auto c_block_dstr = make_static_tile_distribution(c_block_dstr_encode);
return c_block_dstr;
}
};
} // namespace ck_tile