From e34dee57b15de1b3e8a46dafa516ea32bd25b4e9 Mon Sep 17 00:00:00 2001 From: huizzhan Date: Mon, 18 Aug 2025 15:02:16 +0000 Subject: [PATCH] grouped topk ref --- .../ck_tile/09_topk_softmax/topk_softmax.cpp | 21 +++-- .../ck_tile/host/reference/reference_topk.hpp | 90 +++++++++++++++++++ 2 files changed, 102 insertions(+), 9 deletions(-) mode change 100644 => 100755 include/ck_tile/host/reference/reference_topk.hpp diff --git a/example/ck_tile/09_topk_softmax/topk_softmax.cpp b/example/ck_tile/09_topk_softmax/topk_softmax.cpp index 63fd922b3a..acb978a44e 100755 --- a/example/ck_tile/09_topk_softmax/topk_softmax.cpp +++ b/example/ck_tile/09_topk_softmax/topk_softmax.cpp @@ -68,12 +68,12 @@ auto reference_topk_softmax(const ck_tile::HostTensor& x, } template -auto reference_topk_softmax(const ck_tile::HostTensor& x, +void reference_grouped_topk_softmax(const ck_tile::HostTensor& x, ck_tile::HostTensor& y_values, ck_tile::HostTensor& y_indices, - // ck_tile::index_t num_expert_group, - // ck_tile::index_t topk_group, ck_tile::index_t k, + ck_tile::index_t num_expert_group, + ck_tile::index_t topk_group, ck_tile::index_t dim = -1, bool largest = true, bool sorted = true) @@ -81,8 +81,9 @@ auto reference_topk_softmax(const ck_tile::HostTensor& x, using namespace ck_tile; auto y = reference_softmax(x, dim); + printf("==============================before reference_grouped_topk===============================\n"); reference_topk(y, y_values, y_indices, k, dim, largest, sorted); - // reference_grouped_topk(y, y_values, y_indices, k, num_expert_group, topk_group, dim, largest, sorted); + reference_grouped_topk(y, y_values, y_indices, k, num_expert_group, topk_group, dim, largest, sorted); } // different threshold for different dtype @@ -155,8 +156,8 @@ bool test_topk_softmax(ck_tile::ArgParser args) int warmup = args.get_int("warmup"); int repeat = args.get_int("repeat"); - // int num_expert_group = 16; - // int topk_group = 2; + int num_expert_group = 16; + int topk_group = 2; if(stride_input < 0) { @@ -250,9 +251,11 @@ bool test_topk_softmax(ck_tile::ArgParser args) ck_tile::HostTensor value_ref({tokens, topk}, {stride_output, 1}); ck_tile::HostTensor index_ref({tokens, topk}, {stride_output, 1}); - reference_topk_softmax( - x_host, value_ref, index_ref, topk); - // x_host, value_ref, index_ref, num_expert_group, topk_group, topk); + reference_grouped_topk_softmax( + x_host, value_ref, index_ref, topk, num_expert_group, topk_group); + + // reference_grouped_topk_softmax( + // x_host, value_ref, index_ref, topk); auto [rtol, atol] = get_elimit(""); for(int i_t = 0; i_t < tokens; i_t++) 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 0fc99a983a..86434c7d0d --- a/include/ck_tile/host/reference/reference_topk.hpp +++ b/include/ck_tile/host/reference/reference_topk.hpp @@ -122,4 +122,94 @@ CK_TILE_HOST auto reference_topk(const HostTensor& x, return ck_tile::make_tuple(y_values, y_indices); } + + +/* + similiar to vllm grouped_topk() in fused_moe.py + x (Tensor) – the input tensor. + topk (int) – the k in “top-k” + num_expert_group (int) – the number of expert groups + topk_group (int) – the k for expert groups + dim (int, optional) – the dimension to sort along + largest (bool, optional) – largest or smallest elements + sorted (bool, optional) – elements in sorted order or not + output: + y_values + y_indices + https://github.com/ROCm/vllm/blob/main/vllm/model_executor/layers/fused_moe/fused_moe.py#L1657 +*/ +template +CK_TILE_HOST void reference_grouped_topk(const HostTensor& x, + HostTensor& y_values, + HostTensor& y_indices, + index_t topk, + index_t num_expert_group = 16, + index_t topk_group = 2, + index_t dim = -1, + bool largest = true, + bool sorted = true) +{ + printf("==================================reference_grouped_topk===============================\n"); + auto lens = x.get_lengths(); + index_t num_token = lens[0]; + index_t num_expert = lens[1]; + index_t expert_per_group = num_expert / num_expert_group; + + index_t target_dim = (dim == -1) ? (lens.size() - 1) : dim; + assert(target_dim < lens.size()); + assert(k <= lens[target_dim]); + lens[target_dim] = topk; + + HostTensor group_scores({num_token, num_expert_group}, {num_expert_group, 1}); + HostTensor group_mask({num_token, num_expert_group}, {num_expert_group, 1}); + + HostTensor score_mask({num_token, num_expert}, {num_expert, 1}); + HostTensor masked_scores({num_token, num_expert}, {num_expert, 1}); + + HostTensor group_values({num_token, topk_group}, {topk_group, 1}); + HostTensor group_indices({num_token, topk_group}, {topk_group, 1}); + + // calculate group score + auto f1 = [&](auto m) { + for(int n_group = 0; n_group < num_expert_group; ++n_group) { + // max value for expert group + DataType group_max = std::numeric_limits::lowest(); + for(int n = n_group * expert_per_group; n < (n_group + 1) * expert_per_group; ++n) + { + const DataType group_value = x(m, n); + group_max = group_max < group_value ? group_value : group_max; + } + group_scores(m, n_group) = group_max; + } + }; + make_ParallelTensorFunctor(f1, num_token)(std::thread::hardware_concurrency()); + + // select group values and group_indices + reference_topk(group_scores, group_values, group_indices, topk_group, dim, largest, sorted); + + // mask score + auto f2 = [&](auto m) { + // initialize score mask as -inf + for(int n = 0; n < num_expert; ++n) { + score_mask(m, n) = std::numeric_limits::lowest(); + } + // set mask value = 0 for topk groups + for(int k_group = 0; k_group < topk_group; ++k_group) { + int k_group_idx = group_indices(m, k_group); + for(int n = k_group_idx * expert_per_group; n < (k_group_idx + 1) * expert_per_group; ++n) + { + score_mask(m, n) = 0; + } + } + // add mask for scores + for(int n = 0; n < num_expert; ++n) { + masked_scores(m, n) = x(m, n) + score_mask(m, n); + } + }; + make_ParallelTensorFunctor(f2, num_token)(std::thread::hardware_concurrency()); + + // select topk values from masked scores + reference_topk(masked_scores, y_values, y_indices, topk, dim, largest, sorted); + +} } // namespace ck_tile