ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-16 08:44:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

remove debug 9.8 tflops

Browse Source
This commit is contained in:
feifei14119
2025-01-07 15:26:56 +08:00
parent 888317e698
commit fa335f31df
12 changed files with 3746 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
example/ck_tile/18_flatmm_uk/CMakeLists.txt Normal file
View File

@@ -0,0 +1,19 @@
set(TILE_EXAPMLE_FLATMM_UK "tile_example_flatmm_uk")
# not using add_example_executable() to add this target, since we don't want this to have
# to be included in "make all/install/check"
message("adding ${TILE_EXAPMLE_FLATMM_UK}")
file(GLOB INSTANCE_SRCS instances/*.cpp)
add_executable(${TILE_EXAPMLE_FLATMM_UK} EXCLUDE_FROM_ALL main.cpp)
target_include_directories(${TILE_EXAPMLE_FLATMM_UK} PRIVATE ${CMAKE_CURRENT_LIST_DIR})
target_sources(${TILE_EXAPMLE_FLATMM_UK} PRIVATE ${INSTANCE_SRCS})
set(TILE_EXAPMLE_FLATMM_UK_COMPILE_OPTIONS)
# NOTE: we turn off undefined-func-template to let source compile without explicit declare function specializations
list(APPEND TILE_EXAPMLE_FLATMM_UK_COMPILE_OPTIONS -Wno-undefined-func-template -Wno-float-equal)
list(APPEND TILE_EXAPMLE_FLATMM_UK_COMPILE_OPTIONS -DCK_TILE_BUFFER_LOAD_AGPR=1) # TODO: enable load to a
list(APPEND TILE_EXAPMLE_FLATMM_UK_COMPILE_OPTIONS -DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=4) # rta
# list(APPEND TILE_EXAPMLE_FLATMM_UK_COMPILE_OPTIONS -mllvm -greedy-reverse-local-assignment=1)
# list(APPEND TILE_EXAPMLE_FLATMM_UK_COMPILE_OPTIONS -v --save-temps -Wno-gnu-line-marker)
target_compile_options(${TILE_EXAPMLE_FLATMM_UK} PRIVATE ${TILE_EXAPMLE_FLATMM_UK_COMPILE_OPTIONS})

100
example/ck_tile/18_flatmm_uk/flatmm_uk.hpp Normal file
View File

@@ -0,0 +1,100 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck_tile/core.hpp"
#include "ck_tile/host/kernel_launch.hpp"
#include "ck_tile/ops/flatmm_uk.hpp"
#include <string>
// this is only a convenient structure for creating an example
// this is not part of the host API
template <typename I, typename W, typename O, typename ST, typename SW, typename SQ, typename KW>
struct FlatmmUkTypeConfig;
template <typename ST, typename SW, typename SQ, typename KW>
struct FlatmmUkTypeConfig<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, ST, SW, SQ, KW>
{
using ADataType = ck_tile::bf16_t;
using GDataType = ck_tile::bf16_t;
using DDataType = ck_tile::bf16_t;
using AccDataType = float;
using ODataType = ck_tile::bf16_t;
using AScaleDataType = ck_tile::remove_cvref_t<ST>;
using GScaleDataType = ck_tile::remove_cvref_t<SW>;
using DScaleDataType = ck_tile::remove_cvref_t<SW>;
using YSmoothScaleDataType = ck_tile::remove_cvref_t<SQ>;
using TopkWeightDataType = ck_tile::remove_cvref_t<KW>;
using IndexDataType = ck_tile::index_t;
};
template <typename ST, typename SW, typename SQ, typename KW>
struct FlatmmUkTypeConfig<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, ST, SW, SQ, KW>
{
using ADataType = ck_tile::fp16_t;
using GDataType = ck_tile::fp16_t;
using DDataType = ck_tile::fp16_t;
using AccDataType = float;
using ODataType = ck_tile::fp16_t;
using AScaleDataType = ck_tile::remove_cvref_t<ST>;
using GScaleDataType = ck_tile::remove_cvref_t<SW>;
using DScaleDataType = ck_tile::remove_cvref_t<SW>;
using YSmoothScaleDataType = ck_tile::remove_cvref_t<SQ>;
using TopkWeightDataType = ck_tile::remove_cvref_t<KW>;
using IndexDataType = ck_tile::index_t;
};
template <typename ST, typename SW, typename SQ, typename KW>
struct FlatmmUkTypeConfig<ck_tile::int8_t, ck_tile::int8_t, ck_tile::bf16_t, ST, SW, SQ, KW>
{
using ADataType = ck_tile::int8_t;
using GDataType = ck_tile::int8_t;
using DDataType = ck_tile::int8_t;
using AccDataType = int32_t;
using ODataType = ck_tile::bf16_t;
using AScaleDataType = ck_tile::remove_cvref_t<ST>;
using GScaleDataType = ck_tile::remove_cvref_t<SW>;
using DScaleDataType = ck_tile::remove_cvref_t<SW>;
using YSmoothScaleDataType = ck_tile::remove_cvref_t<SQ>;
using TopkWeightDataType = ck_tile::remove_cvref_t<KW>;
using IndexDataType = ck_tile::index_t;
};
struct flatmm_uk_args
{
const void* a_ptr; // [m, k], input token
const void* b_ptr; // [m, k], input token
const void* c_ptr; // [m, k], output token (no need to do zeroing)
void* d_ptr; // [m, k], output token (no need to do zeroing)
void* dbg_int_ptr; // [m, k], output token (no need to do zeroing)
void* dbg_bf16_ptr; // [m, k], output token (no need to do zeroing)
void* dbg_fp32_ptr; // [m, k], output token (no need to do zeroing)
ck_tile::index_t block_m; // block_m, used to devide the input
ck_tile::index_t hidden_size; // k
ck_tile::index_t intermediate_size; // n / TP, for Gate. if Gate+Up, Down need divide by 2
ck_tile::index_t num_tokens; // input number of tokens for current iteration
ck_tile::index_t num_experts; // number of groups
ck_tile::index_t topk; // need this?
ck_tile::index_t stride_token; // for input/output, stride for each row, should >= hidden_size
};
// This is the public API, will be generated by script
struct flatmm_uk_traits
{
std::string prec_i; // input precision
std::string prec_w; // weight precision
std::string prec_o; // output precision
std::string prec_st; // token scale data type
std::string prec_sw; // weight scale data type
std::string prec_sq; // smooth quant scale
std::string prec_kw; // topk-weight data type
int block_m;
int gate_only;
int fused_quant; // 0:no-sweep, 1:smooth-dynamic-quant, 2:dynamic-quant
};
float flatmm_uk(flatmm_uk_traits, flatmm_uk_args, const ck_tile::stream_config&);

192
example/ck_tile/18_flatmm_uk/instances/flatmm_uk_api.cpp Normal file
View File

@@ -0,0 +1,192 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
#include "flatmm_uk.hpp"
#include "flatmm_uk_api.hpp"
#include "ck_tile/ops/flatmm_uk.hpp"
#include <iostream>
template <ck_tile::index_t... Is>
using S = ck_tile::sequence<Is...>;
// do not the define of this tepmlate function inside the _api.cpp, otherwise will block make -j
template <typename Ts_>
float flatmm_uk_(const ck_tile::stream_config& s_, flatmm_uk_args_ a_)
{
printf("[FF] ======= fused_moegemm_() ======= \n \tget moe arg in a_ <flatmm_uk_args>, get "
"config in Ts_\n");
using f_traits = ck_tile::FusedMoeGemmTraits<Ts_::GateOnly, Ts_::FusedQuant == 1, 1 /*atomic*/>;
using f_shape = ck_tile::FusedMoeGemmShape<typename Ts_::BlockTile_0,
typename Ts_::WarpPerBlock_0,
typename Ts_::WarpTile_0,
typename Ts_::BlockTile_1,
typename Ts_::WarpPerBlock_0,
typename Ts_::WarpTile_0>;
printf("[FF] --- fused_moegemm_(): <FusedMoeGemmShape> --- \n");
printf("[FF] f_shape::BlockSize = %d\n", static_cast<uint32_t>(f_shape::BlockSize));
printf("[FF] f_shape::NumWarps = %d\n", static_cast<uint32_t>(f_shape::NumWarps));
printf("[FF] --------- \n");
printf("[FF] f_shape::Block_M0 = %d\n", static_cast<uint32_t>(f_shape::Block_M0));
printf("[FF] f_shape::Block_N0 = %d\n", static_cast<uint32_t>(f_shape::Block_N0));
printf("[FF] f_shape::Block_K0 = %d\n", static_cast<uint32_t>(f_shape::Block_K0));
printf("[FF] f_shape::WarpPerBlock_M0 = %d\n", static_cast<uint32_t>(f_shape::WarpPerBlock_M0));
printf("[FF] f_shape::WarpPerBlock_N0 = %d\n", static_cast<uint32_t>(f_shape::WarpPerBlock_N0));
printf("[FF] f_shape::WarpPerBlock_K0 = %d\n", static_cast<uint32_t>(f_shape::WarpPerBlock_K0));
printf("[FF] f_shape::Warp_M0 = %d\n", static_cast<uint32_t>(f_shape::Warp_M0));
printf("[FF] f_shape::Warp_N0 = %d\n", static_cast<uint32_t>(f_shape::Warp_N0));
printf("[FF] f_shape::Warp_K0 = %d\n", static_cast<uint32_t>(f_shape::Warp_K0));
printf("[FF] f_shape::ThreadPerBlock_M0 = %d\n",
static_cast<uint32_t>(f_shape::ThreadPerBlock_M0));
printf("[FF] f_shape::ThreadPerBlock_N0 = %d\n",
static_cast<uint32_t>(f_shape::ThreadPerBlock_N0));
printf("[FF] f_shape::ThreadPerBlock_K0 = %d\n",
static_cast<uint32_t>(f_shape::ThreadPerBlock_K0));
printf("[FF] f_shape::Repeat_M0 = %d\n", static_cast<uint32_t>(f_shape::Repeat_M0));
printf("[FF] f_shape::Repeat_N0 = %d\n", static_cast<uint32_t>(f_shape::Repeat_N0));
printf("[FF] f_shape::Repeat_K0 = %d\n", static_cast<uint32_t>(f_shape::Repeat_K0));
printf("[FF] f_shape::Block_W0 = %d\n", static_cast<uint32_t>(f_shape::Block_W0));
printf("[FF] f_shape::Block_Nr0 = %d\n", static_cast<uint32_t>(f_shape::Block_Nr0));
printf("[FF] f_shape::Block_Kr0 = %d\n", static_cast<uint32_t>(f_shape::Block_Kr0));
printf("[FF] --------- \n");
printf("[FF] f_shape::Block_M1 = %d\n", static_cast<uint32_t>(f_shape::Block_M1));
printf("[FF] f_shape::Block_N1 = %d\n", static_cast<uint32_t>(f_shape::Block_N1));
printf("[FF] f_shape::Block_K1 = %d\n", static_cast<uint32_t>(f_shape::Block_K1));
printf("[FF] f_shape::WarpPerBlock_M1 = %d\n", static_cast<uint32_t>(f_shape::WarpPerBlock_M1));
printf("[FF] f_shape::WarpPerBlock_N1 = %d\n", static_cast<uint32_t>(f_shape::WarpPerBlock_N1));
printf("[FF] f_shape::WarpPerBlock_K1 = %d\n", static_cast<uint32_t>(f_shape::WarpPerBlock_K1));
printf("[FF] f_shape::Warp_M1 = %d\n", static_cast<uint32_t>(f_shape::Warp_M1));
printf("[FF] f_shape::Warp_N1 = %d\n", static_cast<uint32_t>(f_shape::Warp_N1));
printf("[FF] f_shape::Warp_K1 = %d\n", static_cast<uint32_t>(f_shape::Warp_K1));
printf("[FF] f_shape::ThreadPerBlock_M1 = %d\n",
static_cast<uint32_t>(f_shape::ThreadPerBlock_M1));
printf("[FF] f_shape::ThreadPerBlock_N1 = %d\n",
static_cast<uint32_t>(f_shape::ThreadPerBlock_N1));
printf("[FF] f_shape::ThreadPerBlock_K1 = %d\n",
static_cast<uint32_t>(f_shape::ThreadPerBlock_K1));
printf("[FF] f_shape::Repeat_M1 = %d\n", static_cast<uint32_t>(f_shape::Repeat_M1));
printf("[FF] f_shape::Repeat_N1 = %d\n", static_cast<uint32_t>(f_shape::Repeat_N1));
printf("[FF] f_shape::Repeat_K1 = %d\n", static_cast<uint32_t>(f_shape::Repeat_K1));
printf("[FF] f_shape::Block_W1 = %d\n", static_cast<uint32_t>(f_shape::Block_W1));
printf("[FF] f_shape::Block_Nr1 = %d\n", static_cast<uint32_t>(f_shape::Block_Nr1));
printf("[FF] f_shape::Block_Kr1 = %d\n", static_cast<uint32_t>(f_shape::Block_Kr1));
using f_problem =
ck_tile::FusedMoeGemmPipelineProblem<typename Ts_::ADataType,
typename Ts_::GDataType,
typename Ts_::DDataType,
typename Ts_::AccDataType,
typename Ts_::ODataType,
typename Ts_::AScaleDataType,
typename Ts_::GScaleDataType,
typename Ts_::DScaleDataType,
typename Ts_::YSmoothScaleDataType,
typename Ts_::TopkWeightDataType,
typename Ts_::IndexDataType,
ck_tile::element_wise::FastGeluAsm, // TODO: hardcoded
f_shape,
f_traits>;
// using f_pipeline = ck_tile::FusedMoeGemmPipeline_FlatmmEx<f_problem>;
using f_pipeline = ck_tile::GemmPipeline_FlatmmUk<f_problem>;
using f_kernel = ck_tile::FlatmmUkKernel<f_pipeline, void>;
const dim3 grids = f_kernel::GridSize(a_);
constexpr dim3 blocks = f_kernel::BlockSize();
constexpr ck_tile::index_t kBlockPerCu = 1;
printf("[FF] grids = [%d, %d, %d]\n", grids.x, grids.y, grids.z);
printf("[FF] blocks = [%d, %d, %d]\n", blocks.x, blocks.y, blocks.z);
static int printed = 0;
auto kargs = f_kernel::MakeKargs(a_);
f_kernel kernel{};
auto lambda_kenrel =
ck_tile::make_kernel<blocks.x, kBlockPerCu>(kernel, grids, blocks, 0, kargs);
if(s_.log_level_ > 0 && printed == 10)
{
// std::cout << ", " << f_kernel::GetName() << std::flush;
printed = 1;
}
return ck_tile::launch_kernel(
s_, lambda_kenrel
// ck_tile::make_kernel<blocks.x, kBlockPerCu>(f_kernel{}, grids, blocks, 0, kargs)
);
}
float flatmm_uk(flatmm_uk_traits t, flatmm_uk_args a, const ck_tile::stream_config& s)
{
// auto s_ = ck_tile::stream_config{s.stream_id_, false, s.log_level_, 0, 1};
auto s_ = s;
auto t_ = flatmm_uk_traits_{t.prec_i,
t.prec_w,
t.prec_o,
t.prec_st,
t.prec_sw,
t.prec_sq,
t.prec_kw,
t.block_m,
t.gate_only,
t.fused_quant};
auto a_ = flatmm_uk_args_{
a.a_ptr, // const void* a_ptr;
a.b_ptr, // const void* a_ptr;
a.c_ptr, // void* o_ptr;
a.d_ptr, // void* o_ptr;
a.dbg_int_ptr,
a.dbg_bf16_ptr,
a.dbg_fp32_ptr,
a.hidden_size, // index_t hidden_size;
a.intermediate_size, // index_t intermediate_size;
a.num_tokens, // index_t num_tokens;
a.num_experts, // index_t num_experts;
a.topk, // index_t topk;
a.stride_token // index_t stride_token;
};
float r = -1;
if(t_.prec_i == "bf16" && t_.prec_w == "bf16" && t_.prec_o == "bf16" && t_.prec_st == "fp32" &&
t_.prec_sw == "fp32" && t_.prec_sq == "fp32" && t_.prec_kw == "fp32" && t_.block_m == 32 &&
t_.gate_only == 1)
{
using t_ = fmoe_<ck_tile::bf16_t,
ck_tile::bf16_t,
ck_tile::bf16_t,
float,
float,
float,
float,
S<32, 512, 128, 128>,
S<1, 4, 1>,
S<16, 16, 32>,
1,
0>;
r = flatmm_uk_<t_>(s_, a_);
}
else if(t_.prec_i == "fp16" && t_.prec_w == "fp16" && t_.prec_o == "fp16" &&
t_.prec_st == "fp32" && t_.prec_sw == "fp32" && t_.prec_sq == "fp32" &&
t_.prec_kw == "fp32" && t_.block_m == 32 && t_.gate_only == 1)
{
using t_ = fmoe_<ck_tile::fp16_t,
ck_tile::fp16_t,
ck_tile::fp16_t,
float,
float,
float,
float,
S<32, 512, 128, 128>,
S<1, 4, 1>,
S<16, 16, 32>,
1,
0>;
r = flatmm_uk_<t_>(s_, a_);
}
// keep unsupported case return negative
if(r < 0)
return -1;
return r;
}

76
example/ck_tile/18_flatmm_uk/instances/flatmm_uk_api.hpp Normal file
View File

@@ -0,0 +1,76 @@
// SPDX-License-Identifier: MIT
// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include "ck_tile/core.hpp"
#include "ck_tile/host/kernel_launch.hpp"
#include "ck_tile/ops/flatmm_uk.hpp"
#include <string>
// runtime args
struct flatmm_uk_args_ : public ck_tile::FlatmmUkHostArgs
{
};
// This is the public API, will be generated by script
struct flatmm_uk_traits_
{
std::string prec_i; // input precision
std::string prec_w; // weight precision
std::string prec_o; // output precision
std::string prec_st; // token scale data type
std::string prec_sw; // weight scale data type
std::string prec_sq; // smooth quant scale
std::string prec_kw; // topk-weight data type
int block_m;
int gate_only;
int fused_quant; // 0:no-sweep, 1:smooth-dynamic-quant, 2:dynamic-quant
};
// this is used to pattern-match internl kernel implementation, not to instantiate kernel
template <typename I,
typename W,
typename O,
typename ST,
typename SW,
typename SQ,
typename KW,
typename BlockTIle_, // seq<b_token, b_interm, b_hidden, b_down>
typename WarpPerBlock_,
typename WarpTile_, // seq<*,*,*>, used to select mfma
ck_tile::index_t GateOnly_ = 0,
ck_tile::index_t FusedQuant_ = 0>
struct fmoe_ // traits, ugly name, only used for internal
{
using TypeConfig = FlatmmUkTypeConfig<I, W, O, ST, SW, SQ, KW>;
using ADataType = ck_tile::remove_cvref_t<typename TypeConfig::ADataType>;
using GDataType = ck_tile::remove_cvref_t<typename TypeConfig::GDataType>;
using DDataType = ck_tile::remove_cvref_t<typename TypeConfig::DDataType>;
using AccDataType = ck_tile::remove_cvref_t<typename TypeConfig::AccDataType>;
using ODataType = ck_tile::remove_cvref_t<typename TypeConfig::ODataType>;
using AScaleDataType = ck_tile::remove_cvref_t<typename TypeConfig::AScaleDataType>;
using GScaleDataType = ck_tile::remove_cvref_t<typename TypeConfig::GScaleDataType>;
using DScaleDataType = ck_tile::remove_cvref_t<typename TypeConfig::DScaleDataType>;
using YSmoothScaleDataType = ck_tile::remove_cvref_t<typename TypeConfig::YSmoothScaleDataType>;
using TopkWeightDataType = ck_tile::remove_cvref_t<typename TypeConfig::TopkWeightDataType>;
using IndexDataType = ck_tile::remove_cvref_t<typename TypeConfig::IndexDataType>;
static constexpr ck_tile::index_t BT_ = BlockTIle_::at(ck_tile::number<0>{}); // block token
static constexpr ck_tile::index_t BI_ =
BlockTIle_::at(ck_tile::number<1>{}); // block intermediate
static constexpr ck_tile::index_t BH_ = BlockTIle_::at(ck_tile::number<2>{}); // block hidden
static constexpr ck_tile::index_t BD_ = BlockTIle_::at(ck_tile::number<3>{}); // block down
using BlockTile_0 = ck_tile::sequence<BT_, BI_, BH_>;
using WarpPerBlock_0 = ck_tile::remove_cvref_t<WarpPerBlock_>;
using WarpTile_0 = ck_tile::remove_cvref_t<WarpTile_>;
using BlockTile_1 = ck_tile::sequence<BT_, BD_, BI_ / (GateOnly_ ? 1 : 2)>;
using WarpPerBlock_1 = ck_tile::remove_cvref_t<WarpPerBlock_>;
using WarpTile_1 = ck_tile::remove_cvref_t<WarpTile_>;
static constexpr ck_tile::index_t GateOnly = GateOnly_;
static constexpr ck_tile::index_t FusedQuant = FusedQuant_;
};

692
example/ck_tile/18_flatmm_uk/main.cpp Normal file
View File

@@ -0,0 +1,692 @@
#include <algorithm>
#include <cstring>
#include <unordered_set>
#include <vector>
#include <set>
#include "ck_tile/host.hpp"
#include "flatmm_uk.hpp"
// different threshold for different dtype
template <typename DataType>
auto get_elimit()
{
double rtol = 1e-2;
double atol = 1e-2;
return ck_tile::make_tuple(rtol, atol);
}
template <>
auto get_elimit<ck_tile::bf16_t>()
{
double rtol = 1e-2;
double atol = 1e-2;
return ck_tile::make_tuple(rtol, atol);
}
template <typename ADataType,
typename BDataType,
typename AccDataType,
typename CDataType,
typename AElementOp = ck_tile::identity,
typename BElementOp = ck_tile::identity,
typename ACCElementOp = ck_tile::identity>
CK_TILE_HOST void my_reference_gemm(const ck_tile::HostTensor<ADataType>& a_m_k,
const ck_tile::HostTensor<BDataType>& b_k_n,
ck_tile::HostTensor<CDataType>& c_m_n,
float t,
const AElementOp& a_element_op = {},
const BElementOp& b_element_op = {},
const ACCElementOp& acc_element_op = {})
{
const std::size_t M = a_m_k.get_length(0);
const std::size_t N = b_k_n.get_length(0);
const std::size_t K = a_m_k.get_length(1);
printf("[REF] M = %zu, N = %zu, K = %zu\n", M, N, K);
auto cal_tflops = [&](auto ms) {
double flop_gemm = 2.0 * M * N * K;
return (flop_gemm) / (static_cast<double>(ms) * 1e-3) / 1e12;
};
auto cal_tbps = [&](auto ms) {
double a_bytes = static_cast<double>(M) * K * sizeof(ADataType);
double b_bytes = static_cast<double>(N) * K * sizeof(BDataType);
double o_bytes = static_cast<double>(M) * N * sizeof(CDataType);
return (a_bytes + b_bytes + o_bytes) / (static_cast<double>(ms) * 1e-3) / 1e12;
};
std::cout << ", " << t * 1.E3 << " us, " << cal_tflops(t) << " tflops, " << cal_tbps(t)
<< " TB/s" << std::endl
<< std::flush;
auto f_mn = [&](auto m, auto n) {
AccDataType v_acc = 0;
for(std::size_t k = 0; k < K; ++k)
{
ADataType v_a = a_element_op(a_m_k(m, k));
BDataType v_b = b_element_op(b_k_n(n, k));
v_acc +=
ck_tile::type_convert<AccDataType>(v_a) * ck_tile::type_convert<AccDataType>(v_b);
}
c_m_n(m, n) = ck_tile::type_convert<CDataType>(acc_element_op(v_acc));
};
ck_tile::make_ParallelTensorFunctor(f_mn, M, N)(std::thread::hardware_concurrency());
}
// mfma_type, 0:32x32, 1:16x16
// TODO: padding?
template <typename T>
auto shuffle_moe_weight(const ck_tile::HostTensor<T>& t, std::string mfma_dtype, int mfma_type = 0)
{
assert(t.get_lengths().size() == 3);
int b_ = t.get_lengths()[0];
int n_ = t.get_lengths()[1];
int k_ = t.get_lengths()[2];
if((mfma_dtype == "bf16" || mfma_dtype == "fp16") && mfma_type == 0)
{
ck_tile::HostTensor<T> t_view({b_, n_ / 32, 32, k_ / 16, 2, 8});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 1, 3, 4, 2, 5});
}
else if((mfma_dtype == "bf16" || mfma_dtype == "fp16") && mfma_type == 1)
{
ck_tile::HostTensor<T> t_view({b_, n_ / 16, 16, k_ / 32, 4, 8});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 1, 3, 4, 2, 5});
}
else if((mfma_dtype == "int8" || mfma_dtype == "fp8") && mfma_type == 0)
{
ck_tile::HostTensor<T> t_view({b_, n_ / 32, 32, k_ / 32, 2, 16});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 1, 3, 4, 2, 5});
}
else if((mfma_dtype == "int8" || mfma_dtype == "fp8") && mfma_type == 1)
{
ck_tile::HostTensor<T> t_view({b_, n_ / 16, 16, k_ / 64, 4, 16});
std::copy(t.begin(), t.end(), t_view.begin());
return ck_tile::reference_permute(t_view, {0, 1, 3, 4, 2, 5});
}
return t;
}
template <typename T>
auto shuffle_weight(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()[0];
int k_ = t.get_lengths()[1];
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});
printf("[FF] permute: n_ = %d, k_ = %d, n_/16 = %d, k_/32 = %d\n", n_, k_, n_ / 16, k_ / 32);
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;
}
auto create_args(int argc, char* argv[])
{
ck_tile::ArgParser arg_parser;
arg_parser.insert("m", "64", "num of m")
.insert("n", "1024", "num of n")
.insert("k", "8192", "num of k")
.insert("t", "64", "num input tokens")
.insert("e", "8", "num of experts")
.insert("tk", "1", "topk")
.insert("h", "4096", "hidden_size of this model")
.insert("i", "4096", "intermediate_size between 2 gemms of FFN")
.insert("stride", "-1", "stride per row, if -1 then equal to hidden_size")
.insert("bm", "32", "blocking factor for sorted tokens")
.insert("tp", "8", "tensor parallel size")
.insert("v", "1", "cpu validation or not")
.insert("kname", "1", "print kernel name or not")
.insert("prec_i", "bf16", "input precision")
.insert("prec_w", "bf16", "weight precision")
.insert("prec_o", "bf16", "output precision")
.insert("prec_st", "auto", "token scale data type. auto will set to fp32")
.insert("prec_sw", "auto", "weight scale data type. auto will set to fp32")
.insert("prec_sq", "auto", "(dynamic) smooth quant data type. auto will set to fp32")
.insert("prec_kw", "auto", "topk-weight data type. auto will set to fp32")
.insert("fquant", "0", "fused-quant, 0:no, 1:smooth-dynamic-quant, 2:dynamic-quant")
.insert(
"gate_only", "1", "w0(gate/up) style, 0:gate+up will double interm size, 1:only gate")
.insert("api", "0", "benchmark api set: 0:fused-moe(moe-gemm+moe-sorting), 1:moe-gemm")
.insert("balance",
"0",
"if set to 1, will try balance the expert in topk-ids(convenient for testing)")
.insert("init",
"2",
"init method. 0:random stepped float(fast). 1: random uniform, 2:rand normalized"
"normalized(slow)")
.insert("seed", "11939", "seed used to do random")
.insert("warmup", "1", "cold iter")
.insert("repeat", "4", "hot iter");
bool result = arg_parser.parse(argc, argv);
return std::make_tuple(result, arg_parser);
}
// I:input-type, W:weight-type, O:output-type, ST:toke-scale-tpye, SW:weight-scale-type,
// SQ:smooth-quant-type, KW:topk-weight-type
template <typename I, typename W, typename O, typename ST, typename SW, typename SQ, typename KW>
bool run(const ck_tile::ArgParser& arg_parser)
{
ck_tile::index_t M = arg_parser.get_int("m");
ck_tile::index_t N = arg_parser.get_int("n");
ck_tile::index_t K = arg_parser.get_int("k");
printf("[FF] M = %d, N = %d, K = %d\n", M, N, K);
ck_tile::index_t experts = arg_parser.get_int("e");
ck_tile::index_t topk = arg_parser.get_int("tk");
ck_tile::index_t stride = arg_parser.get_int("stride");
ck_tile::index_t block_m = arg_parser.get_int("bm");
std::string prec_i = arg_parser.get_str("prec_i");
std::string prec_w = arg_parser.get_str("prec_w");
std::string prec_o = arg_parser.get_str("prec_o");
std::string prec_st = arg_parser.get_str("prec_st");
std::string prec_sw = arg_parser.get_str("prec_sw");
std::string prec_sq = arg_parser.get_str("prec_sq");
std::string prec_kw = arg_parser.get_str("prec_kw");
prec_st = (prec_st == "auto") ? "fp32" : prec_st;
prec_sw = (prec_sw == "auto") ? "fp32" : prec_sw;
prec_sq = (prec_sq == "auto") ? "fp32" : prec_sq;
prec_kw = (prec_kw == "auto") ? "fp32" : prec_kw;
int kname = arg_parser.get_int("kname");
int do_validation = arg_parser.get_int("v");
int warmup = arg_parser.get_int("warmup");
int repeat = arg_parser.get_int("repeat");
int fused_quant = arg_parser.get_int("fquant");
int gate_only = arg_parser.get_int("gate_only");
int init = arg_parser.get_int("init");
uint32_t seed = arg_parser.get_uint32("seed");
using TypeConfig = FlatmmUkTypeConfig<I, W, O, ST, SW, SQ, KW>;
using ADataType = typename TypeConfig::ADataType;
using BDataType = ADataType;
using AccDataType = typename TypeConfig::AccDataType;
using CDataType = AccDataType;
using DDataType = AccDataType;
// host verify
ck_tile::HostTensor<ADataType> a_host({M, K});
ck_tile::HostTensor<BDataType> b_host({N, K});
ck_tile::HostTensor<CDataType> c_host({M, N});
ck_tile::HostTensor<DDataType> d_host({M, N});
ck_tile::HostTensor<int> dbg_int({M * N, K});
ck_tile::HostTensor<float> dbg_fp32({M * N, K});
ck_tile::HostTensor<ck_tile::bf16_t> dbg_bf16({M * N, K});
if(init == 0)
{
ck_tile::FillStepRange<ADataType>{-.5f, .5f, 0.01f}(a_host);
ck_tile::FillStepRange<BDataType>{-.5f, .5f, 0.01f}(b_host);
}
else if(init == 1)
{
ck_tile::FillUniformDistribution<ADataType>{-.5f, .5f, seed, true}(a_host);
ck_tile::FillUniformDistribution<BDataType>{-.5f, .5f, seed, true}(b_host);
}
else if(init == 2)
{
ck_tile::FillNormalDistribution<ADataType>{0.f, 1.f, seed, true}(a_host);
ck_tile::FillNormalDistribution<BDataType>{0.f, 1.f, seed, true}(b_host);
}
/*
// a_host
{
int X = static_cast<int>(K);
int Y = static_cast<int>(M);
for(int y = 0; y < Y; y++)
{
for(int x = 0; x < X; x++)
{
int idx = X * y + x;
a_host.mData[idx] = ck_tile::type_convert<ADataType>(x * 1.0f);
//b_host.mData[idx] = ck_tile::type_convert<GDataType>(y * 1.0f);
//b_host.mData[idx] = ck_tile::type_convert<GDataType>(y*1.f + x * 0.0001f);
}
}
}
// b_host
{
int X = static_cast<int>(K);
int Y = static_cast<int>(N);
for(int y = 0; y < Y; y++)
{
for(int x = 0; x < X; x++)
{
int idx = X * y + x;
b_host.mData[idx] = ck_tile::type_convert<GDataType>(idx * 1.0f);
//b_host.mData[idx] = ck_tile::type_convert<GDataType>(y * 1.0f);
//b_host.mData[idx] = ck_tile::type_convert<GDataType>(y*1.f + x * 0.0001f);
}
}
}*/
// permute weight
ck_tile::HostTensor<BDataType> b_perm_host = shuffle_weight(b_host, prec_w, 1);
ck_tile::DeviceMem a_buf(a_host);
ck_tile::DeviceMem b_buf(b_perm_host); // b_host -> b_perm_host
ck_tile::DeviceMem c_buf(c_host);
ck_tile::DeviceMem d_buf(d_host);
ck_tile::DeviceMem dbg_int_buf(dbg_int);
ck_tile::DeviceMem dbg_bf16_buf(dbg_bf16);
ck_tile::DeviceMem dbg_fp32_buf(dbg_fp32);
flatmm_uk_traits traits{prec_i,
prec_w,
prec_o,
prec_st,
prec_sw,
prec_sq,
prec_kw,
block_m,
gate_only,
fused_quant};
printf("[FF] --- run(): <flatmm_uk_traits> ---\n");
printf("[FF] traits.prec_i = %s\n", traits.prec_i.c_str());
printf("[FF] traits.prec_w = %s\n", traits.prec_w.c_str());
printf("[FF] traits.prec_o = %s\n", traits.prec_o.c_str());
printf("[FF] traits.prec_st = %s\n", traits.prec_st.c_str());
printf("[FF] traits.prec_sw = %s\n", traits.prec_sw.c_str());
printf("[FF] traits.prec_sq = %s\n", traits.prec_sq.c_str());
printf("[FF] traits.prec_kw = %s\n", traits.prec_kw.c_str());
printf("[FF] traits.block_m = %d\n", traits.block_m);
printf("[FF] traits.gate_only = %d\n", traits.gate_only);
printf("[FF] traits.fused_quant = %d\n", traits.fused_quant);
flatmm_uk_args args{a_buf.GetDeviceBuffer(),
b_buf.GetDeviceBuffer(),
c_buf.GetDeviceBuffer(),
d_buf.GetDeviceBuffer(),
dbg_int_buf.GetDeviceBuffer(),
dbg_bf16_buf.GetDeviceBuffer(),
dbg_fp32_buf.GetDeviceBuffer(),
block_m,
K,
N,
M,
experts,
topk,
stride};
printf("[FF] --- run(): <flatmm_uk_args> ---\n");
printf("[FF] args.block_m = %d\n", args.block_m);
printf("[FF] args.hidden_size = %d\n", args.hidden_size);
printf("[FF] args.intermediate_size = %d\n", args.intermediate_size);
printf("[FF] args.num_tokens = %d\n", args.num_tokens); // 1
printf("[FF] args.topk = %d\n", args.topk); // 0
printf("[FF] args.num_experts = %d\n", args.num_experts); // 0
printf("[FF] args.stride_token = %d\n", args.stride_token);
float ave_time = flatmm_uk(
traits, args, ck_tile::stream_config{nullptr, true, kname ? 1 : 0, warmup, repeat});
if(ave_time < 0)
{
std::cout << " not supported!" << std::endl << std::flush;
return false;
}
bool pass = true;
if(do_validation)
{
auto d_dev = d_buf.ToHost<float>();
std::cout << std::endl << " =================== " << std::endl;
d_host.SetZero();
my_reference_gemm<ADataType, BDataType, CDataType, DDataType>(
a_host, b_host, d_host, ave_time);
pass = ck_tile::check_err(d_dev, d_host);
std::cout << "The CPU veification result is:" << (pass ? "correct" : "fail") << std::endl;
}
#if 0
int GridDimX = 2;
int GridDimY = 1;
int BlockDimX = 64;
int BlockDimY = 4;
int BlockSize = BlockDimX * BlockDimY;
// dbg_int
{
auto dbg_int_dev = dbg_int_buf.ToHost<int>();
std::ofstream file("ff_dbg_int.txt");
file << " [dbg_int]: Grid = [" << GridDimX << ", " << GridDimY << "], Block = " << BlockSize
<< std::endl;
for(int bidy = 0; bidy < GridDimY; bidy++)
{
for(int bidx = 0; bidx < GridDimX; bidx++)
{
file << "\n ========== block : [" << bidx << ", " << bidy << "] ==========";
for(int tid = 0; tid < BlockSize; tid++)
{
int gid = (BlockSize * GridDimX) * bidy + BlockSize * bidx + tid;
if(tid % 64 == 0)
{
file << "\n [" << tid << " : " << tid + 63 << "]: ";
}
file << ck_tile::type_convert<int>(dbg_int_dev.mData[gid]) << ", ";
}
}
}
file.close();
}
// dbg_bf16 ---> kernel
{
auto dbg_bf16_dev = dbg_bf16_buf.ToHost<BDataType>();
std::ofstream file("ff_dbg_bf16_kernel.txt");
file << " [dbg_bf16]: Grid = [" << GridDimX << ", " << GridDimY
<< "], Block = " << BlockSize << std::endl;
for(int bidy = 0; bidy < GridDimY; bidy++)
{
for(int bidx = 0; bidx < GridDimX; bidx++)
{
file << "\n ========== block : [" << bidx << ", " << bidy << "] ==========";
for(int tid = 0; tid < BlockSize; tid++)
{
int gid = (BlockSize * bidx) * bidy + BlockSize * bidx + tid;
file << "\n [" << tid << "]: ";
for(int i = 0; i < 64; i++) // multi output per thread
file << ck_tile::type_convert<float>(dbg_bf16_dev.mData[gid * 64 + i])
<< ", ";
}
}
}
file.close();
}
// dbg_bf16
{
auto dbg_bf16_dev = dbg_bf16_buf.ToHost<BDataType>();
std::ofstream file("ff_dbg_bf16.txt");
int X = static_cast<int>(N);
int Y = static_cast<int>(M);
file << " [dbg_bf16]: Row = " << Y << ", Col = " << X << std::endl;
for(int m = 0; m < Y; m++)
{
file << "\n ========== row : [" << m << " / " << Y << "] ==========";
for(int n = 0; n < X; n++)
{
if(n % 64 == 0)
{
file << "\n [" << n << " : " << n + 63 << "]: ";
}
int idx = X * m + n;
file << ck_tile::type_convert<float>(dbg_bf16_dev.mData[idx]) << ", ";
}
}
file.close();
}
// dbg_fp32 ---> kernel
{
auto dbg_fp32_dev = dbg_fp32_buf.ToHost<float>();
std::ofstream file("ff_dbg_fp32_kernel.txt");
file << " [dbg_fp32]: Grid = [" << GridDimX << ", " << GridDimY
<< "], Block = " << BlockSize << std::endl;
for(int bidy = 0; bidy < GridDimY; bidy++)
{
for(int bidx = 0; bidx < GridDimX; bidx++)
{
file << "\n ========== block : [" << bidx << ", " << bidy << "] ==========";
for(int tid = 0; tid < BlockSize; tid++)
{
int gid = (BlockSize * bidx) * bidy + BlockSize * bidx + tid;
file << "\n [" << tid << "]: ";
for(int i = 0; i < 64; i++) // multi output per thread
file << ck_tile::type_convert<float>(dbg_fp32_dev.mData[gid * 64 + i])
<< ", ";
// if(tid % 64 == 0) // one output per thread
// file << "\n [" << tid << " : " << tid + 63 << "]: ";
// file << ck_tile::type_convert<float>(dbg_bf16.mData[gid]) << ", ";
}
}
}
file.close();
}
// dbg_fp32
{
auto dbg_fp32_dev = dbg_fp32_buf.ToHost<float>();
std::ofstream file("ff_dbg_fp32.txt");
int X = static_cast<int>(N);
int Y = static_cast<int>(M);
file << " [dbg_fp32]: Row = " << Y << ", Col = " << X << std::endl;
for(int m = 0; m < Y; m++)
{
file << "\n ========== row : [" << m << " / " << Y << "] ==========";
for(int n = 0; n < X; n++)
{
if(n % 64 == 0)
{
file << "\n [" << n << " : " << n + 63 << "]: ";
}
int idx = X * m + n;
file << ck_tile::type_convert<float>(dbg_fp32_dev.mData[idx]) << ", ";
}
}
file.close();
}
// a_host
{
std::ofstream file("ff_a_host.txt");
int X = static_cast<int>(K);
int Y = static_cast<int>(M);
file << " [a_host]: Row = " << Y << ", Col = " << X << std::endl;
for(int y = 0; y < Y; y++)
{
file << "\n ========== row : [" << y << " / " << Y << "] ==========";
for(int x = 0; x < X; x++)
{
int idx = X * y + x;
if(idx % 16 == 0)
{
file << "\n [" << x << " : " << x + 15 << " ]: ";
}
file << ck_tile::type_convert<float>(a_host.mData[idx]) << ", ";
}
}
file.close();
}
// b_host
{
std::ofstream file("ff_b_host.txt");
int X = static_cast<int>(K);
int Y = static_cast<int>(N);
file << " [b_host]: Row = " << Y << ", Col = " << X << std::endl;
for(int y = 0; y < Y; y++)
{
file << "\n ========== row : [" << y << " / " << Y << "] ==========";
for(int x = 0; x < X; x++)
{
int idx = X * y + x;
if(idx % 16 == 0)
{
file << "\n [" << x << " : " << x + 15 << " ]: ";
}
file << ck_tile::type_convert<float>(b_host.mData[idx]) << ", ";
}
}
file.close();
}
// permute_b
{
std::ofstream file("ff_b_perm_host.txt");
int X = static_cast<int>(K);
int Y = static_cast<int>(N);
file << " [b_perm_host]: Row = " << Y << ", Col = " << X << std::endl;
for(int y = 0; y < Y; y++)
{
file << "\n ========== row : [" << y << " / " << Y << "] ==========";
for(int x = 0; x < X; x++)
{
int idx = X * y + x;
if(idx % 16 == 0)
{
file << "\n [" << x << " : " << x + 15 << " ]: ";
}
file << ck_tile::type_convert<float>(b_perm_host.mData[idx]) << ", ";
}
}
file.close();
}
// d_dev ---> kernel
{
auto d_dev = d_buf.ToHost<float>();
std::ofstream file("ff_d_dev_kernel.txt");
file << " [d_dev]: Grid = [" << GridDimX << ", " << GridDimY << "], Block = " << BlockSize
<< std::endl;
for(int bidy = 0; bidy < GridDimY; bidy++)
{
for(int bidx = 0; bidx < GridDimX; bidx++)
{
file << "\n ========== block : [" << bidx << ", " << bidy << "] ==========";
for(int tid = 0; tid < BlockSize; tid++)
{
int gid = (BlockSize * bidx) * bidy + BlockSize * bidx + tid;
file << "\n [" << tid << "]: ";
for(int i = 0; i < 64; i++) // multi output per thread
file << ck_tile::type_convert<float>(d_dev.mData[gid * 64 + i]) << ", ";
}
}
}
file.close();
}
// d_dev
{
auto d_dev = d_buf.ToHost<float>();
std::ofstream file("ff_d_dev.txt");
int X = static_cast<int>(N);
int Y = static_cast<int>(M);
file << " [d_dev]: Row = " << Y << ", Col = " << X << std::endl;
for(int y = 0; y < Y; y++)
{
file << "\n ========== row : [" << y << " / " << Y << "] ==========";
for(int x = 0; x < X; x++)
{
if(x % 64 == 0)
{
file << "\n [" << x << " : " << x + 63 << "]: ";
}
int idx = X * y + x;
file << ck_tile::type_convert<float>(d_dev.mData[idx]) << ", ";
}
}
file.close();
}
// d_host
{
std::ofstream file("ff_d_host.txt");
int X = static_cast<int>(N);
int Y = static_cast<int>(M);
file << " [d_host]: Row = " << Y << ", Col = " << X << std::endl;
for(int y = 0; y < Y; y++)
{
file << "\n ========== row : [" << y << " / " << Y << "] ==========";
for(int x = 0; x < X; x++)
{
if(x % 64 == 0)
{
file << "\n [" << x << " : " << x + 63 << "]: ";
}
int idx = X * y + x;
file << ck_tile::type_convert<float>(d_host.mData[idx]) << ", ";
}
}
file.close();
}
#endif
std::cout << std::flush << std::endl;
return pass;
}
int main(int argc, char* argv[])
{
auto [result, arg_parser] = create_args(argc, argv);
if(!result)
return -1;
std::string prec_i = arg_parser.get_str("prec_i");
std::string prec_w = arg_parser.get_str("prec_w");
std::string prec_o = arg_parser.get_str("prec_o");
std::string prec_st = arg_parser.get_str("prec_st");
std::string prec_sw = arg_parser.get_str("prec_sw");
std::string prec_sq = arg_parser.get_str("prec_sq");
std::string prec_kw = arg_parser.get_str("prec_kw");
prec_st = (prec_st == "auto") ? "fp32" : prec_st;
prec_sw = (prec_sw == "auto") ? "fp32" : prec_sw;
prec_sq = (prec_sq == "auto") ? "fp32" : prec_sq;
prec_kw = (prec_kw == "auto") ? "fp32" : prec_kw;
// no dynamic quant case
if(prec_i == "bf16" && prec_w == "bf16" && prec_o == "bf16" && prec_kw == "fp32")
{
return run<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float>(
arg_parser)
? 0
: -2;
}
else if(prec_i == "fp16" && prec_w == "fp16" && prec_o == "fp16" && prec_kw == "fp32")
{
return run<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float>(
arg_parser)
? 0
: -2;
}
return -3;
}

1
example/ck_tile/CMakeLists.txt
View File

@@ -17,3 +17,4 @@ add_subdirectory(14_moe_smoothquant)
add_subdirectory(15_fused_moe)
add_subdirectory(16_batched_gemm)
add_subdirectory(17_grouped_gemm)
add_subdirectory(18_flatmm_uk)