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

grouped topk ref

Browse Source
This commit is contained in:
huizzhan
2025-08-18 15:02:16 +00:00
parent bb918bfd04
commit e34dee57b1
2 changed files with 102 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
example/ck_tile/09_topk_softmax/topk_softmax.cpp
View File

@@ -68,12 +68,12 @@ auto reference_topk_softmax(const ck_tile::HostTensor<InputType>& x,
}
template <typename InputType, typename WeightType, typename IndexType = ck_tile::index_t>
auto reference_topk_softmax(const ck_tile::HostTensor<InputType>& x,
void reference_grouped_topk_softmax(const ck_tile::HostTensor<InputType>& x,
ck_tile::HostTensor<WeightType>& y_values,
ck_tile::HostTensor<IndexType>& 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<InputType>& x,
using namespace ck_tile;
auto y = reference_softmax<InputType, WeightType, WeightType>(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<WeightType> value_ref({tokens, topk}, {stride_output, 1});
ck_tile::HostTensor<IndexType> index_ref({tokens, topk}, {stride_output, 1});
reference_topk_softmax<InputType, WeightType, IndexType>(
x_host, value_ref, index_ref, topk);
// x_host, value_ref, index_ref, num_expert_group, topk_group, topk);
reference_grouped_topk_softmax<InputType, WeightType, IndexType>(
x_host, value_ref, index_ref, topk, num_expert_group, topk_group);
// reference_grouped_topk_softmax<InputType, WeightType, IndexType>(
// x_host, value_ref, index_ref, topk);
auto [rtol, atol] = get_elimit<InputType>("");
for(int i_t = 0; i_t < tokens; i_t++)