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

topk acc fixed with gemm random input

Browse Source
This commit is contained in:
huizzhan
2025-08-12 11:06:44 +00:00
parent e0248c67f2
commit 7b85491177
3 changed files with 98 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -476,24 +476,25 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
// apply topk for softmax output
auto x_tmp = p_compute;
// // initialize x_tmp for topk debug
// printf("===============on device debug input=====================\n");
// std::mt19937 rng(123);
// std::uniform_int_distribution<int> dist_debug_input(1, 100);
constexpr auto x_tmp_spans = decltype(x_tmp)::get_distributed_spans();
sweep_tile_span(x_tmp_spans[number<0>{}], [&](auto idx0) {
sweep_tile_span(x_tmp_spans[number<1>{}], [&](auto idx1) {
const auto tile_idx = get_x_indices_from_distributed_indices(
x_tmp.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);
x_tmp(i_j_idx) = sin(float(row_id + col_id)) * 100;
// x_tmp(i_j_idx) = sin(float(row_id)) * cos(float(col_id));
// x_tmp(i_j_idx) = float(dist_debug_input(rng));
});
});
// // // initialize x_tmp for topk debug
// // printf("===============on device debug input=====================\n");
// // std::mt19937 rng(123);
// // std::uniform_int_distribution<int> dist_debug_input(1, 100);
// constexpr auto x_tmp_spans = decltype(x_tmp)::get_distributed_spans();
// sweep_tile_span(x_tmp_spans[number<0>{}], [&](auto idx0) {
// sweep_tile_span(x_tmp_spans[number<1>{}], [&](auto idx1) {
// const auto tile_idx = get_x_indices_from_distributed_indices(
// x_tmp.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);
// x_tmp(i_j_idx) = sin(float(row_id + col_id)) * 100;
// // x_tmp(i_j_idx) = sin(float(row_id)) * cos(float(col_id));
// // x_tmp(i_j_idx) = float(dist_debug_input(rng));
// });
// });
// argmax for topk
const auto f_argmax = [](ArgmaxPacket e0, ArgmaxPacket e1) {
@@ -535,12 +536,8 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
auto col_id = tile_idx.at(number<1>{});
constexpr auto i_j_idx = make_tuple(idx0, idx1);
ArgmaxPacket tmp = r(i_idx);
// if (col_id == i_k) {
// // debug_block_tile(i_j_idx) = tmp.value;
// debug_block_tile(i_j_idx) = tmp.arg;
// }
debug_block_tile(i_j_idx) = (col_id == i_k) ? tmp.value: debug_block_tile(i_j_idx);
// debug_block_tile(i_j_idx) = (col_id == i_k) ? tmp.arg: debug_block_tile(i_j_idx);
// debug_block_tile(i_j_idx) = (col_id == i_k) ? tmp.value: debug_block_tile(i_j_idx);
debug_block_tile(i_j_idx) = (col_id == i_k) ? tmp.arg: debug_block_tile(i_j_idx);
// value_block_tile(i_j_idx) = tmp.value;
// index_block_tile(i_j_idx) = tmp.arg;
});
@@ -561,7 +558,6 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg
});
}
return debug_block_tile;
// return x_tmp;
}
};

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

@@ -70,7 +70,8 @@ int main(int argc, char* argv[])
using CDataType = ck_tile::half_t;
using WeightType = float;
using DebugType = float;
using IndexType = ck_tile::index_t;
// using IndexType = ck_tile::index_t;
using IndexType = float;
ck_tile::index_t verification = 0;
ck_tile::index_t M = 3328;
@@ -157,32 +158,35 @@ int main(int argc, char* argv[])
ck_tile::HostTensor<WeightType> value_host_dev(out_lengths, out_strides);
ck_tile::HostTensor<IndexType> index_host_dev(out_lengths, out_strides);
ck_tile::FillUniformDistributionIntegerValue<ADataType>{-5.f, 5.f}(a_host);
ck_tile::FillUniformDistributionIntegerValue<BDataType>{-5.f, 5.f}(b_host);
// ck_tile::FillUniformDistributionIntegerValue<ADataType>{-5.f, 5.f}(a_host);
// ck_tile::FillUniformDistributionIntegerValue<BDataType>{-5.f, 5.f}(b_host);
ck_tile::HostTensor<WeightType> debug_host_input(debug_lengths, debug_strides);
ck_tile::FillUniformDistribution<ADataType>{0.01f, 0.05f}(a_host);
ck_tile::FillUniformDistribution<BDataType>{0.01f, 0.05f}(b_host);
// ck_tile::HostTensor<WeightType> debug_host_input(debug_lengths, debug_strides);
ck_tile::HostTensor<DebugType> debug_host_dev(debug_lengths, debug_strides);
// std::random_device rd;
// std::mt19937 gen(rd());
// std::uniform_real_distribution<> dist_b(0.01f, 0.05f);
printf("===============debug input=====================\n");
// std::mt19937 rng(123);
// std::uniform_int_distribution<int> dist_debug_input(1, 100);
for(int m = 0; m < M; ++m) {
printf("m: %d ", m);
for(int n = 0; n < N; ++n) {
// debug_host_input(m, n) = float(dist_debug_input(rng));
debug_host_input(m, n) = sin(float(m + n)) * 100;
printf("[%d]:%.4f ", n, debug_host_input(m, n));
}
printf("/n");
}
// // std::random_device rd;
// // std::mt19937 gen(rd());
// // std::uniform_real_distribution<> dist_b(0.01f, 0.05f);
// printf("===============debug input=====================\n");
// // std::mt19937 rng(123);
// // std::uniform_int_distribution<int> dist_debug_input(1, 100);
// for(int m = 0; m < M; ++m) {
// printf("m: %d ", m);
// for(int n = 0; n < N; ++n) {
// // debug_host_input(m, n) = float(dist_debug_input(rng));
// debug_host_input(m, n) = sin(float(m + n)) * 100;
// printf("[%d]:%.4f ", n, debug_host_input(m, n));
// }
// printf("/n");
// }
// std::random_device rd;
// std::mt19937 gen(rd());
// std::uniform_real_distribution<> dist_a(0.f, 0.01f);
// std::uniform_real_distribution<> dist_b(0.01f, 0.05f);
// // std::random_device rd;
// std::mt19937 gen(12345);
// std::uniform_real_distribution<> dist_a(0.f, 0.001f);
// std::uniform_real_distribution<> dist_b(0.001f, 0.005f);
// for(int m = 0; m < M; ++m) {
// for(int k = 0; k < K; ++k) {
@@ -196,11 +200,11 @@ int main(int argc, char* argv[])
// }
// }
// for(std::size_t i = 0; i < 20; ++i) {
// const double a = *std::next(std::begin(a_host), i);
// const double b = *std::next(std::begin(b_host), i);
// std::cout << " a[" << i << "]: " << a << " " << "b[" << i << "]: " << b << std::endl;
// }
for(std::size_t i = 0; i < 20; ++i) {
const double a = *std::next(std::begin(a_host), i);
const double b = *std::next(std::begin(b_host), i);
std::cout << " a[" << i << "]: " << a << " " << "b[" << i << "]: " << b << std::endl;
}
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());
@@ -279,10 +283,12 @@ int main(int argc, char* argv[])
ck_tile::HostTensor<DebugType> 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_topk(debug_host_input, 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);
// reference_topk(debug_host_input, value_ref, index_ref, topk);
// debug_ref = reference_basic_gemm<ADataType, ADataType, AccDataType>(a_host, b_host);
// debug_ref = reference_basic_gemm_softmax<ADataType, ADataType, AccDataType>(a_host, b_host);
reference_basic_gemm_softmax_grouped_topk<ADataType, ADataType, AccDataType, WeightType, IndexType>(
a_host, b_host, value_ref, index_ref, topk);
debug_buf.FromDevice(debug_host_dev.mData.data());
value_buf.FromDevice(value_host_dev.mData.data());
index_buf.FromDevice(index_host_dev.mData.data());
@@ -302,8 +308,12 @@ int main(int argc, char* argv[])
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_end =
// std::vector<size_t>{static_cast<size_t>(i_t + 1), static_cast<size_t>(N)};
auto s_debug_host = debug_host_dev.slice(s_begin, s_end);
auto s_debug_ref = value_ref.slice(s_begin, s_end);
// auto s_debug_ref = debug_ref.slice(s_begin, s_end);
// auto s_debug_ref = value_ref.slice(s_begin, s_end);
auto s_debug_ref = index_ref.slice(s_begin, s_end);
rtn &= ck_tile::check_err(s_debug_host,
s_debug_ref,
std::string("[") + std::to_string(i_t) +
@@ -314,14 +324,14 @@ int main(int argc, char* argv[])
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]:%.4f ", i, r);
printf("ref[%zu]:%.8f ", 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]:%.4f ", i, o);
printf("out[%zu]:%.8f ", i, o);
// std::cout << i_t << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl;
}
printf("\n");

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

@@ -125,3 +125,37 @@ auto reference_basic_gemm_softmax(const ck_tile::HostTensor<ADataType>& a_m_k,
// reference_grouped_topk(c_m_n, y_values, y_indices, topk, num_expert_group, topk_group, dim, largest, sorted);
return c_m_n;
}
template <typename ADataType, typename BDataType, typename AccDataType>
auto reference_basic_gemm(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);
}
};
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;
}