From 71cea3dd708e833da0fdaddd3b7f38c52f04acf7 Mon Sep 17 00:00:00 2001 From: huizzhan Date: Wed, 13 Aug 2025 10:07:29 +0000 Subject: [PATCH] grouped topk acc fixed, little error for index --- ...grouped_topk_pipeline_agmem_bgmem_creg.hpp | 210 ++++++++++++------ .../gemm_softmax_grouped_topk.cpp | 16 +- .../reference_gemm_softmax_grouped_topk.hpp | 4 +- .../ck_tile/host/reference/reference_topk.hpp | 2 +- 4 files changed, 150 insertions(+), 82 deletions(-) mode change 100644 => 100755 include/ck_tile/host/reference/reference_topk.hpp diff --git a/example/ck_tile/39_gemm_softmax_grouped_topk/block_gemm_softmax_grouped_topk_pipeline_agmem_bgmem_creg.hpp b/example/ck_tile/39_gemm_softmax_grouped_topk/block_gemm_softmax_grouped_topk_pipeline_agmem_bgmem_creg.hpp index 5c3aa6b55d..b5a684e6dc 100755 --- a/example/ck_tile/39_gemm_softmax_grouped_topk/block_gemm_softmax_grouped_topk_pipeline_agmem_bgmem_creg.hpp +++ b/example/ck_tile/39_gemm_softmax_grouped_topk/block_gemm_softmax_grouped_topk_pipeline_agmem_bgmem_creg.hpp @@ -441,8 +441,8 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg auto p_compute = make_static_distributed_tensor(c_block_tile.get_tile_distribution()); - // auto debug_block_tile = - // make_static_distributed_tensor(p_compute.get_tile_distribution()); + auto debug_block_tile = + make_static_distributed_tensor(p_compute.get_tile_distribution()); constexpr auto p_spans = decltype(p_compute)::get_distributed_spans(); @@ -470,7 +470,6 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg constexpr auto i_j_idx = make_tuple(idx0, idx1); p_compute(i_j_idx) = p_compute[i_j_idx] / rowsum_p[i_idx]; - // debug_block_tile(i_j_idx) = p_compute(i_j_idx); }); }); @@ -478,7 +477,7 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg auto x_tmp = p_compute; // calculate group score, need to creat group scores tensor int num_expert_group = 16; - // int topk_group = 2; + int topk_group = 2; int expert_per_group = kNPerBlock / num_expert_group; constexpr auto p_compute_spans = decltype(p_compute)::get_distributed_spans(); auto group_scores = x_tmp; @@ -525,87 +524,156 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg // x_tmp_3d, sequence<2>{}, f_max, std::numeric_limits::lowest()); // block_tile_reduce_sync(group_scores, f_max); - // // Step2: select group values and group_indices - // // argmax for topk - // const auto f_argmax = [](ArgmaxPacket e0, ArgmaxPacket e1) { - // return e0.value > e1.value ? e0 : e1; - // }; - // auto group_packet = topk(group_scores, topk_group) + // Step2: select topk group and cal mask score matrix + // argmax for topk + const auto f_argmax = [](ArgmaxPacket e0, ArgmaxPacket e1) { + return e0.value > e1.value ? e0 : e1; + }; - // // Step3: mask score matrix + // // topk_group_index = topk(group_scores, topk_group) + // auto topk_group_index = x_tmp; + // // init topk_group_index to -inf // 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 != group_packet(i_j_idx).arg) ? -numeric::infinity() - // : x_tmp(i_j_idx); + // topk_group_index(i_j_idx) = -numeric::infinity(); // }); // }); - // // Step4: select topk values from masked scores - // for(index_t i_k = 0; i_k < topk; i_k++) - // { - // constexpr auto p_compute_spans = decltype(p_compute)::get_distributed_spans(); - // auto packet = [&]() { - // auto tmp = make_static_distributed_tensor(p_compute.get_tile_distribution()); - // 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.value = x_tmp(i_j_idx); // !!! we reference p_compute here - // t.arg = tile_idx.at(number<1>{}); - // tmp(i_j_idx) = t; - // }); - // }); - // return tmp; - // }(); + // topk_group_mask(1 for selected group scores, -inf for other group scores) + auto topk_group_scores_mask = x_tmp; + // init topk_group_scores_mask to -inf + sweep_tile_span(p_compute_spans[number<0>{}], [&](auto idx0) { + sweep_tile_span(p_compute_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + topk_group_scores_mask(i_j_idx) = -numeric::infinity(); + }); + }); - // auto argmax_init = ArgmaxPacket{-numeric::infinity(), 0}; - // auto r = block_tile_reduce(packet, sequence<1>{}, f_argmax, argmax_init); + for(index_t k_group = 0; k_group < topk_group; k_group++) + { + auto group_packet = [&]() { + auto tmp = make_static_distributed_tensor(p_compute.get_tile_distribution()); + 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.value = group_scores(i_j_idx); // !!! we reference p_compute here + t.arg = tile_idx.at(number<1>{}); + tmp(i_j_idx) = t; + }); + }); + return tmp; + }(); - // block_tile_reduce_xor_sync(r, f_argmax); + auto argmax_init = ArgmaxPacket{-numeric::infinity(), 0}; + auto group_r = block_tile_reduce(group_packet, sequence<1>{}, f_argmax, argmax_init); - // // constexpr auto value_spans = decltype(value_block_tile)::get_distributed_spans(); + block_tile_reduce_xor_sync(group_r, f_argmax); + + // 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_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; - // }); - // }); + 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); + auto k_group_idx = group_r(i_idx).arg; + // topk_group_index(i_j_idx) = (col_id == k_group) ? tmp.value: topk_group_index(i_j_idx); + // topk_group_index(i_j_idx) = (col_id == k_group) ? k_group_idx: topk_group_index(i_j_idx); + topk_group_scores_mask(i_j_idx) = ((col_id >= (k_group_idx * expert_per_group)) && (col_id < ((k_group_idx + 1) * expert_per_group))) ? 1 : topk_group_scores_mask(i_j_idx); + }); + }); - // // update value - // 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 col_id = tile_idx.at(number<1>{}); + // update value + 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 col_id = tile_idx.at(number<1>{}); - // constexpr auto i_j_idx = make_tuple(idx0, idx1); + constexpr auto i_j_idx = make_tuple(idx0, idx1); - // x_tmp(i_j_idx) = (col_id == r(i_idx).arg) ? -numeric::infinity() - // : x_tmp(i_j_idx); - // }); - // }); - // } - // return debug_block_tile; - return group_scores; + group_scores(i_j_idx) = (col_id == group_r(i_idx).arg) ? -numeric::infinity() + : group_scores(i_j_idx); + }); + }); + } + + // Step3: mask score matrix + auto x_tmp_masked = x_tmp; + sweep_tile_span(p_compute_spans[number<0>{}], [&](auto idx0) { + sweep_tile_span(p_compute_spans[number<1>{}], [&](auto idx1) { + constexpr auto i_j_idx = make_tuple(idx0, idx1); + x_tmp_masked(i_j_idx) = x_tmp(i_j_idx) * topk_group_scores_mask(i_j_idx); + }); + }); + + // Step4: select topk values from masked score matrix + for(index_t i_k = 0; i_k < topk; i_k++) + { + auto packet = [&]() { + auto tmp = make_static_distributed_tensor(p_compute.get_tile_distribution()); + 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.value = x_tmp_masked(i_j_idx); // !!! we reference p_compute here + t.arg = tile_idx.at(number<1>{}); + tmp(i_j_idx) = t; + }); + }); + return tmp; + }(); + + auto argmax_init = ArgmaxPacket{-numeric::infinity(), 0}; + auto r = block_tile_reduce(packet, sequence<1>{}, f_argmax, argmax_init); + + block_tile_reduce_xor_sync(r, f_argmax); + + // 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_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; + }); + }); + + // update value + 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 col_id = tile_idx.at(number<1>{}); + + constexpr auto i_j_idx = make_tuple(idx0, idx1); + + x_tmp_masked(i_j_idx) = (col_id == r(i_idx).arg) ? -numeric::infinity() + : x_tmp_masked(i_j_idx); + }); + }); + } + return debug_block_tile; + // return x_tmp_masked; } }; diff --git a/example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.cpp b/example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.cpp index 80b8dcd114..0a810f1f05 100755 --- a/example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.cpp +++ b/example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.cpp @@ -286,9 +286,9 @@ int main(int argc, char* argv[]) // reference_topk(debug_host_input, value_ref, index_ref, topk); // debug_ref = reference_basic_gemm(a_host, b_host); - debug_ref = reference_basic_gemm_softmax(a_host, b_host); - // reference_basic_gemm_softmax_grouped_topk( - // a_host, b_host, value_ref, index_ref, topk); + // debug_ref = reference_basic_gemm_softmax(a_host, b_host); + reference_basic_gemm_softmax_grouped_topk( + 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()); @@ -306,14 +306,14 @@ int main(int argc, char* argv[]) for(int i_t = 0; i_t < tokens; i_t++) { auto s_begin = std::vector{static_cast(i_t), static_cast(0)}; - // auto s_end = - // std::vector{static_cast(i_t + 1), static_cast(topk)}; auto s_end = - std::vector{static_cast(i_t + 1), static_cast(N)}; + std::vector{static_cast(i_t + 1), static_cast(topk)}; + // auto s_end = + // std::vector{static_cast(i_t + 1), static_cast(N)}; auto s_debug_host = debug_host_dev.slice(s_begin, s_end); - auto s_debug_ref = debug_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); + 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) + diff --git a/example/ck_tile/39_gemm_softmax_grouped_topk/reference_gemm_softmax_grouped_topk.hpp b/example/ck_tile/39_gemm_softmax_grouped_topk/reference_gemm_softmax_grouped_topk.hpp index e99d81cf33..ff9187cf7e 100755 --- a/example/ck_tile/39_gemm_softmax_grouped_topk/reference_gemm_softmax_grouped_topk.hpp +++ b/example/ck_tile/39_gemm_softmax_grouped_topk/reference_gemm_softmax_grouped_topk.hpp @@ -63,8 +63,8 @@ void reference_basic_gemm_softmax_grouped_topk(const ck_tile::HostTensor diff --git a/include/ck_tile/host/reference/reference_topk.hpp b/include/ck_tile/host/reference/reference_topk.hpp old mode 100644 new mode 100755 index a20ba93fa6..f7309098db --- a/include/ck_tile/host/reference/reference_topk.hpp +++ b/include/ck_tile/host/reference/reference_topk.hpp @@ -145,7 +145,7 @@ CK_TILE_HOST void reference_grouped_topk(const HostTensor& x, HostTensor& y_values, HostTensor& y_indices, index_t topk, - index_t num_expert_group = 4, + index_t num_expert_group = 16, index_t topk_group = 2, index_t dim = -1, bool largest = true,