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

topk debug3

Browse Source
This commit is contained in:
huizzhan
2025-08-11 10:54:12 +00:00
parent 78dd9a1d89
commit f7ce8aa4f7
5 changed files with 218 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files

102
example/ck_tile/39_gemm_softmax_grouped_topk/block_gemm_softmax_grouped_topk_pipeline_agmem_bgmem_creg.hpp
View File

@@ -38,8 +38,8 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
// for topk computing
struct ArgmaxPacket
{
WeightType arg;
IndexType value;
WeightType value;
IndexType arg;
};
using BlockGemm = remove_cvref_t<decltype(Policy::template GetBlockGemm<Problem>())>;
@@ -208,13 +208,15 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
}
#endif
template <typename ADramBlockWindowTmp, typename BDramBlockWindowTmp, typename ValueBlockTile, typename IndexBlockTile>
CK_TILE_HOST_DEVICE void operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp,
// template <typename ADramBlockWindowTmp, typename BDramBlockWindowTmp, typename DebugBlockTile, typename ValueBlockTile, typename IndexBlockTile>
template <typename ADramBlockWindowTmp, typename BDramBlockWindowTmp>
CK_TILE_HOST_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp,
const BDramBlockWindowTmp& b_dram_block_window_tmp,
ValueBlockTile& value_block_tile,
IndexBlockTile& index_block_tile,
index_t num_loop,
void* p_smem) const
// DebugBlockTile& debug_block_tile,
// ValueBlockTile& value_block_tile,
// IndexBlockTile& index_block_tile,
{
static_assert(
std::is_same_v<ADataType, remove_cvref_t<typename ADramBlockWindowTmp::DataType>> &&
@@ -437,9 +439,8 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
auto p_compute =
make_static_distributed_tensor<ComputeDataType>(c_block_tile.get_tile_distribution());
auto wgid = blockIdx.x + blockIdx.y * gridDim.x + gridDim.x * gridDim.y * blockIdx.z;
auto tid = (threadIdx.z * (blockDim.x * blockDim.y)) + (threadIdx.y * blockDim.x) + threadIdx.x;
int i = 0, j = 0;
auto debug_block_tile =
make_static_distributed_tensor<WeightType>(p_compute.get_tile_distribution());
constexpr auto p_spans = decltype(p_compute)::get_distributed_spans();
@@ -450,13 +451,8 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
constexpr auto i_j_idx = make_tuple(idx0, idx1);
p_compute(i_j_idx) = exp(c_block_tile[i_j_idx] - m_local[i_idx]);
if (wgid == 0 && tid ==0 && i == 0) {
// if (i == 0 && j == 0) {
printf("============c_block_tile(i %d, j %d): %f===============\n", i, j, c_block_tile(i_j_idx));
}
j++;
});
i++;
});
// rowsum for p_compute{i, j}
@@ -466,21 +462,14 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
block_tile_reduce_sync(rowsum_p, f_sum);
// softmax = p_compute{i, j} / rowsum_p
i = 0;
j = 0;
sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) {
constexpr auto i_idx = make_tuple(idx0);
sweep_tile_span(p_spans[number<1>{}], [&](auto idx1) {
constexpr auto i_j_idx = make_tuple(idx0, idx1);
p_compute(i_j_idx) = p_compute[i_j_idx] / rowsum_p[i_idx];
if (wgid == 0 && tid ==0 && i == 0) {
// if (i == 0 && j == 0) {
printf("============p_compute(i %d, j %d): %f===============\n", i, j, p_compute(i_j_idx));
}
j++;
// debug_block_tile(i_j_idx) = p_compute(i_j_idx);
});
i++;
});
// apply topk for softmax output
@@ -489,81 +478,66 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
// argmax for topk
const auto f_argmax = [](ArgmaxPacket e0, ArgmaxPacket e1) {
return e0.arg > e1.arg ? e0 : e1;
return e0.value > e1.value ? e0 : e1;
};
printf("==========================topk: %d====================================\n", topk);
for(index_t i_k = 0; i_k < topk; i_k++)
for(index_t i_k = 0; i_k < 1; i_k++)
{
printf("==========================i_k: %d====================================\n", i_k);
constexpr auto span_2d = decltype(p_compute)::get_distributed_spans();
constexpr auto p_compute_spans = decltype(p_compute)::get_distributed_spans();
auto packet = [&]() {
auto tmp = make_static_distributed_tensor<ArgmaxPacket>(p_compute.get_tile_distribution());
i = 0;
j = 0;
sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) {
sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) {
sweep_tile_span(p_compute_spans[number<0>{}], [&](auto idx0) {
sweep_tile_span(p_compute_spans[number<1>{}], [&](auto idx1) {
const auto tile_idx = get_x_indices_from_distributed_indices(
tmp.get_tile_distribution(), make_tuple(idx0, idx1));
constexpr auto i_j_idx = make_tuple(idx0, idx1);
ArgmaxPacket t;
t.arg = x_tmp(i_j_idx); // !!! we reference p_compute here
t.value = tile_idx.at(number<1>{});
t.value = x_tmp(i_j_idx); // !!! we reference p_compute here
t.arg = tile_idx.at(number<1>{});
tmp(i_j_idx) = t;
// if (wgid == 0 && tid ==0 && i == 0) {
// printf("=========p_compute(i %d, j %d)- t.arg: %f t.arg: %d=======\n", i, j, t.arg, t.value);
// }
j++;
});
i++;
});
return tmp;
}();
i = 0;
j = 0;
sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) {
sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) {
constexpr auto i_j_idx = make_tuple(idx0, idx1);
if (wgid == 0 && tid ==0 && i == 0) {
printf("====packet(i %d, j %d)- t.arg: %f t.arg: %d====\n", i, j, packet(i_j_idx).arg, packet(i_j_idx).value);
}
j++;
});
i++;
});
auto argmax_init = ArgmaxPacket{-numeric<WeightType>::infinity(), 0};
auto r = block_tile_reduce<ArgmaxPacket>(packet, sequence<1>{}, f_argmax, argmax_init);
block_tile_reduce_xor_sync(r, f_argmax);
// auto value_block_tile = make_static_distributed_tensor<WeightType>(dst_dist);
// auto index_block_tile = make_static_distributed_tensor<IndexType>(dst_dist);
sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) {
sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) {
// constexpr auto value_spans = decltype(value_block_tile)::get_distributed_spans();
sweep_tile_span(p_compute_spans[number<0>{}], [&](auto idx0) {
constexpr auto i_idx = make_tuple(idx0);
sweep_tile_span(p_compute_spans[number<1>{}], [&](auto idx1) {
const auto tile_idx = get_x_indices_from_distributed_indices(
p_compute.get_tile_distribution(), make_tuple(idx0, idx1));
// auto row_id = tile_idx.at(number<0>{});
auto col_id = tile_idx.at(number<1>{});
constexpr auto i_j_idx = make_tuple(idx0, idx1);
ArgmaxPacket tmp = r(i_j_idx);
value_block_tile(i_j_idx) = tmp.arg;
index_block_tile(i_j_idx) = tmp.value;
ArgmaxPacket tmp = r(i_idx);
// debug_block_tile(i_j_idx) = (col_id == i_k) ? tmp.value: 0;
debug_block_tile(i_j_idx) = (col_id == i_k) ? tmp.arg: 0;
// value_block_tile(i_j_idx) = tmp.value;
// index_block_tile(i_j_idx) = tmp.arg;
});
});
// update value
sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) {
sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) {
sweep_tile_span(p_compute_spans[number<0>{}], [&](auto idx0) {
sweep_tile_span(p_compute_spans[number<1>{}], [&](auto idx1) {
const auto tile_idx = get_x_indices_from_distributed_indices(
p_compute.get_tile_distribution(), make_tuple(idx0, idx1));
auto col_id = tile_idx.at(number<1>{});
constexpr auto i_j_idx = make_tuple(idx0, idx1);
x_tmp(i_j_idx) = (col_id == r(i_j_idx).value) ? -numeric<WeightType>::infinity()
x_tmp(i_j_idx) = (col_id == r(i_j_idx).arg) ? -numeric<WeightType>::infinity()
: x_tmp(i_j_idx);
});
});
}
return debug_block_tile;
}
};

120
example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.cpp
View File

@@ -74,7 +74,7 @@ int main(int argc, char* argv[])
ck_tile::index_t M = 3328;
ck_tile::index_t N = 4096;
ck_tile::index_t K = 4096;
ck_tile::index_t topk = 16;
ck_tile::index_t topk = 8;
if(argc == 2)
{
@@ -143,8 +143,8 @@ int main(int argc, char* argv[])
const auto b_lengths = std::array<ck_tile::index_t, 2>{N, K};
const auto b_strides = std::array<ck_tile::index_t, 2>{Ldb, 1};
// const auto c_lengths = std::array<ck_tile::index_t, 2>{M, N};
// const auto c_strides = std::array<ck_tile::index_t, 2>{Ldc, 1};
const auto debug_lengths = std::array<ck_tile::index_t, 2>{M, N};
const auto debug_strides = std::array<ck_tile::index_t, 2>{N, 1};
const auto out_lengths = std::array<ck_tile::index_t, 2>{M, topk};
const auto out_strides = std::array<ck_tile::index_t, 2>{Ldout, 1};
@@ -152,7 +152,7 @@ int main(int argc, char* argv[])
// host verify
ck_tile::HostTensor<ADataType> a_host(a_lengths, a_strides);
ck_tile::HostTensor<BDataType> b_host(b_lengths, b_strides);
// ck_tile::HostTensor<CDataType> c_host_dev(c_lengths, c_strides);
ck_tile::HostTensor<WeightType> debug_host_dev(debug_lengths, debug_strides);
ck_tile::HostTensor<WeightType> value_host_dev(out_lengths, out_strides);
ck_tile::HostTensor<IndexType> index_host_dev(out_lengths, out_strides);
@@ -184,7 +184,7 @@ int main(int argc, char* argv[])
ck_tile::DeviceMem a_buf(a_host.get_element_space_size_in_bytes());
ck_tile::DeviceMem b_buf(b_host.get_element_space_size_in_bytes());
// ck_tile::DeviceMem c_buf(c_host_dev.get_element_space_size_in_bytes());
ck_tile::DeviceMem debug_buf(debug_host_dev.get_element_space_size_in_bytes());
ck_tile::DeviceMem value_buf(value_host_dev.get_element_space_size_in_bytes());
ck_tile::DeviceMem index_buf(index_host_dev.get_element_space_size_in_bytes());
@@ -207,7 +207,7 @@ int main(int argc, char* argv[])
constexpr ck_tile::index_t kGemmKPerBlock = 16;
#endif
constexpr ck_tile::index_t kGemmNPerBlock = 256;
constexpr ck_tile::index_t kGemmTopKPerBlock = 16;
constexpr ck_tile::index_t kGemmTopKPerBlock = 8;
ck_tile::index_t kGridSize = (M / kGemmMPerBlock) * (N / kGemmNPerBlock);
@@ -241,6 +241,7 @@ int main(int argc, char* argv[])
0,
static_cast<ADataType*>(a_buf.GetDeviceBuffer()),
static_cast<BDataType*>(b_buf.GetDeviceBuffer()),
static_cast<WeightType*>(debug_buf.GetDeviceBuffer()),
static_cast<WeightType*>(value_buf.GetDeviceBuffer()),
static_cast<IndexType*>(index_buf.GetDeviceBuffer()),
M,
@@ -255,61 +256,96 @@ int main(int argc, char* argv[])
bool rtn = true;
if(verification)
{
ck_tile::HostTensor<WeightType> debug_ref({M, N}, {N, 1});
ck_tile::HostTensor<WeightType> value_ref(out_lengths, out_strides);
ck_tile::HostTensor<IndexType> index_ref(out_lengths, out_strides);
reference_basic_gemm_softmax_grouped_topk<ADataType, ADataType, AccDataType, WeightType, IndexType>(
a_host, b_host, value_ref, index_ref, topk);
// reference_basic_gemm_softmax_grouped_topk<ADataType, ADataType, AccDataType, WeightType, IndexType>(
// a_host, b_host, value_ref, index_ref, topk);
debug_ref = reference_basic_gemm_softmax<ADataType, ADataType, AccDataType>(a_host, b_host);
debug_buf.FromDevice(debug_host_dev.mData.data());
value_buf.FromDevice(value_host_dev.mData.data());
index_buf.FromDevice(index_host_dev.mData.data());
// pass &= ck_tile::check_err(c_host_dev, c_host_ref);
// rtn &= ck_tile::check_err(debug_host_dev, debug_ref);
// for(std::size_t i = 0; i < debug_ref.size(); ++i) {
// const double o = *std::next(std::begin(debug_host_dev), i);
// const double r = *std::next(std::begin(debug_ref), i);
// std::cout << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl;
// }
// std::cout << "valid:" << (rtn ? "y" : "n") << std::endl;
const ck_tile::index_t tokens = M;
auto [rtol, atol] = get_elimit<ADataType>("");
auto [rtol, atol] = get_elimit<WeightType>("");
for(int i_t = 0; i_t < tokens; i_t++)
{
auto s_begin = std::vector<size_t>{static_cast<size_t>(i_t), static_cast<size_t>(0)};
auto s_end =
std::vector<size_t>{static_cast<size_t>(i_t + 1), static_cast<size_t>(topk)};
auto s_value_host = value_host_dev.slice(s_begin, s_end);
auto s_value_ref = value_ref.slice(s_begin, s_end);
rtn &= ck_tile::check_err(s_value_host,
s_value_ref,
auto s_debug_host = debug_host_dev.slice(s_begin, s_end);
auto s_debug_ref = debug_ref.slice(s_begin, s_end);
rtn &= ck_tile::check_err(s_debug_host,
s_debug_ref,
std::string("[") + std::to_string(i_t) +
std::string("] Value Error:"),
rtol,
atol);
// for(std::size_t i = 0; i < s_value_ref.size(); ++i) {
// const double o = *std::next(std::begin(s_value_host), i);
// const double r = *std::next(std::begin(s_value_ref), i);
// std::cout << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl;
// }
auto s_index_host = index_host_dev.slice(s_begin, s_end);
auto s_index_ref = index_ref.slice(s_begin, s_end);
rtn &= ck_tile::check_err(s_index_host,
s_index_ref,
std::string("[") + std::to_string(i_t) +
std::string("] Index Error:"),
rtol,
atol);
// for(std::size_t i = 0; i < s_index_ref.size(); ++i) {
// const double o = *std::next(std::begin(s_index_host), i);
// const double r = *std::next(std::begin(s_index_ref), i);
// std::cout << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl;
// }
printf("row [%d]\n", i_t);
for(std::size_t i = 0; i < s_debug_ref.size(); ++i) {
// const double o = *std::next(std::begin(s_debug_host), i);
const double r = *std::next(std::begin(s_debug_ref), i);
printf("ref[%zu]:%f ", i, r);
// std::cout << i_t << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl;
}
printf("\n");
for(std::size_t i = 0; i < s_debug_ref.size(); ++i) {
const double o = *std::next(std::begin(s_debug_host), i);
// const double r = *std::next(std::begin(s_debug_ref), i);
printf("out[%zu]:%f ", i, o);
// std::cout << i_t << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl;
}
printf("\n");
}
std::cout << "valid:" << (rtn ? "y" : "n") << std::endl;
}
// const ck_tile::index_t tokens = M;
// auto [rtol, atol] = get_elimit<ADataType>("");
// for(int i_t = 0; i_t < tokens; i_t++)
// {
// auto s_begin = std::vector<size_t>{static_cast<size_t>(i_t), static_cast<size_t>(0)};
// auto s_end =
// std::vector<size_t>{static_cast<size_t>(i_t + 1), static_cast<size_t>(topk)};
// auto s_value_host = value_host_dev.slice(s_begin, s_end);
// auto s_value_ref = value_ref.slice(s_begin, s_end);
// rtn &= ck_tile::check_err(s_value_host,
// s_value_ref,
// std::string("[") + std::to_string(i_t) +
// std::string("] Value Error:"),
// rtol,
// atol);
// // for(std::size_t i = 0; i < s_value_ref.size(); ++i) {
// // const double o = *std::next(std::begin(s_value_host), i);
// // const double r = *std::next(std::begin(s_value_ref), i);
// // std::cout << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl;
// // }
// auto s_index_host = index_host_dev.slice(s_begin, s_end);
// auto s_index_ref = index_ref.slice(s_begin, s_end);
// rtn &= ck_tile::check_err(s_index_host,
// s_index_ref,
// std::string("[") + std::to_string(i_t) +
// std::string("] Index Error:"),
// rtol,
// atol);
// // for(std::size_t i = 0; i < s_index_ref.size(); ++i) {
// // const double o = *std::next(std::begin(s_index_host), i);
// // const double r = *std::next(std::begin(s_index_ref), i);
// // std::cout << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl;
// // }
// }
// std::cout << "valid:" << (rtn ? "y" : "n") << std::endl;
// }
std::size_t flop = std::size_t(2) * M * N * K;
std::size_t num_btype =
sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N;
float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
float gb_per_sec = num_btype / 1.E6 / ave_time;
std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
<< std::endl;
std::cout << "Perf: " << ave_time << " ms, " << std::endl;
return rtn;
// return !rtn;
}

8
example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.hpp
View File

@@ -179,6 +179,7 @@ struct Gemm
CK_TILE_DEVICE void operator()(const ADataType* p_a,
const BDataType* p_b,
WeightType* p_debug,
WeightType* p_value,
IndexType* p_index,
const index_t M,
@@ -200,6 +201,11 @@ struct Gemm
p_b, make_tuple(N, K), make_tuple(Ldb, 1), number<kBAlignment>{}, number<1>{});
}();
const auto debug_dram = [&] {
return make_naive_tensor_view<address_space_enum::global>(
p_debug, make_tuple(M, N), make_tuple(N, 1), number<kOutAlignment>{}, number<1>{});
}();
const auto value_dram = [&] {
return make_naive_tensor_view<address_space_enum::global>(
p_value, make_tuple(M, topK), make_tuple(Ldout, 1), number<kOutAlignment>{}, number<1>{});
@@ -210,7 +216,7 @@ struct Gemm
p_index, make_tuple(M, topK), make_tuple(Ldout, 1), number<kOutAlignment>{}, number<1>{});
}();
GridGemm{}(a_dram, b_dram, value_dram, index_dram, c_element_func);
GridGemm{}(a_dram, b_dram, debug_dram, value_dram, index_dram, c_element_func);
}
};

39
example/ck_tile/39_gemm_softmax_grouped_topk/grid_gemm_softmax_grouped_topk.hpp
View File

@@ -23,9 +23,10 @@ struct GridGemm
static constexpr auto topk = Policy::kTopKPerBlock;
static constexpr auto kBlockSize = Policy::kBlockSize;
template <typename AGridTensorView, typename BGridTensorView, typename ValueGridTensorView, typename IndexGridTensorView>
template <typename AGridTensorView, typename BGridTensorView, typename DebugGridTensorView, typename ValueGridTensorView, typename IndexGridTensorView>
CK_TILE_DEVICE void operator()(const AGridTensorView& a_grid,
const BGridTensorView& b_grid,
DebugGridTensorView& debug_grid,
ValueGridTensorView& value_grid,
IndexGridTensorView& index_grid,
const CElementFunction& c_element_func) const
@@ -70,7 +71,7 @@ struct GridGemm
// constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM;
// constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK;
constexpr index_t K1 = 16 / sizeof( WeightType);
constexpr index_t K1 = 16 / sizeof(WeightType);
constexpr index_t K0 = topk / K1;
constexpr index_t M2 = get_warp_size() / K0;
// coalesce reading for each blocks
@@ -109,9 +110,39 @@ struct GridGemm
tile_elementwise_inout([](auto& value) { value = 0; }, value_block_tile);
tile_elementwise_inout([](auto& index) { index = 0; }, index_block_tile);
block_gemm_pipeline(a_block_window, b_block_window, value_block_tile, index_block_tile, K / kKPerBlock, p_smem_char);
// constexpr index_t debugK1 = 16 / sizeof(WeightType);
// constexpr index_t debugK0 = kNPerBlock / debugK1;
// constexpr index_t debugM2 = get_warp_size() / debugK0;
// // coalesce reading for each blocks
// constexpr index_t debugM1 = kBlockSize / get_warp_size();
// constexpr index_t debugM0 = kMPerBlock / (debugM2 * debugM1);
auto debug_window = make_tile_window(
debug_grid, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {iM, iN});
// auto debug_window = make_tile_window(
// debug_grid, make_tuple(number<kMPerBlock>{}, number<kNPerBlock>{}), {iM, iN},
// make_static_tile_distribution(
// tile_distribution_encoding<sequence<1>,
// tuple<sequence<debugM0, debugM1, debugM2>, sequence<debugK0, debugK1>>,
// tuple<sequence<1>, sequence<1, 2>>,
// tuple<sequence<1>, sequence<2, 0>>,
// sequence<1, 2>,
// sequence<0, 1>>{}));
// using DebugBlockTileDistr = decltype(debug_window.get_tile_distribution());
// using DebugBlockTile = decltype(make_static_distributed_tensor<WeightType>(DebugBlockTileDistr{}));
// DebugBlockTile debug_block_tile;
// tile_elementwise_inout([](auto& debug) { debug = 0; }, debug_block_tile);
// block_gemm_pipeline(a_block_window, b_block_window, debug_block_tile, value_block_tile, index_block_tile, K / kKPerBlock, p_smem_char);
const auto debug_block_tile = block_gemm_pipeline(a_block_window, b_block_window, K / kKPerBlock, p_smem_char);
// cast DataType and apply CElementFunction
const auto debug_cast_block_tile = tile_elementwise_in(
[&](const auto& debug) { return c_element_func(type_convert<WeightType>(debug)); },
debug_block_tile);
const auto value_cast_block_tile = tile_elementwise_in(
[&](const auto& value) { return c_element_func(type_convert<WeightType>(value)); },
value_block_tile);
@@ -120,6 +151,8 @@ struct GridGemm
[&](const auto& index) { return c_element_func(type_convert<IndexType>(index)); },
index_block_tile);
store_tile(debug_window, debug_cast_block_tile);
store_tile(value_window, value_cast_block_tile);
store_tile(index_window, index_cast_block_tile);
}

59
example/ck_tile/39_gemm_softmax_grouped_topk/reference_gemm_softmax_grouped_topk.hpp
View File

@@ -66,3 +66,62 @@ void reference_basic_gemm_softmax_grouped_topk(const ck_tile::HostTensor<ADataTy
reference_topk(c_m_n, y_values, y_indices, topk);
// reference_grouped_topk(c_m_n, y_values, y_indices, topk, num_expert_group, topk_group, dim, largest, sorted);
}
template <typename ADataType, typename BDataType, typename AccDataType>
auto reference_basic_gemm_softmax(const ck_tile::HostTensor<ADataType>& a_m_k,
const ck_tile::HostTensor<BDataType>& b_n_k)
{
const int M = a_m_k.mDesc.get_lengths()[0];
const int N = b_n_k.mDesc.get_lengths()[0];
const int K = b_n_k.mDesc.get_lengths()[1];
ck_tile::HostTensor<AccDataType> c_m_n({M, N}, {N, 1});
auto f = [&](auto m) {
for(int n = 0; n < N; ++n)
{
AccDataType v_acc = 0;
for(int k = 0; k < K; ++k)
{
ADataType v_a = a_m_k(m, k);
BDataType v_b = b_n_k(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<AccDataType>(v_acc);
}
// reference softmax
AccDataType v_max = std::numeric_limits<ADataType>::lowest();
// max
for(int n = 0; n < N; ++n)
{
const AccDataType v_c = c_m_n(m, n);
v_max = v_max < v_c ? v_c : v_max;
}
AccDataType v_exp_sum = 0;
// sum
for(int n = 0; n < N; ++n)
{
const AccDataType v_c = c_m_n(m, n);
v_exp_sum += ck_tile::exp(v_c - v_max);
}
// elementwise
for(int n = 0; n < N; ++n)
{
const AccDataType v_c = c_m_n(m, n);
c_m_n(m, n) = ck_tile::exp(v_c - v_max) / v_exp_sum;
}
};
ck_tile::make_ParallelTensorFunctor(f, c_m_n.mDesc.get_lengths()[0])(
std::thread::hardware_concurrency());
// reference_topk(c_m_n, y_values, y_indices, topk);
// reference_grouped_topk(c_m_n, y_values, y_indices, topk, num_expert_group, topk_group, dim, largest, sorted);
return c_m_n;
}