ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-12 02:05:50 +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++)

90
include/ck_tile/host/reference/reference_topk.hpp Normal file → Executable file
View File

@@ -122,4 +122,94 @@ CK_TILE_HOST auto reference_topk(const HostTensor<DataType>& 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 <typename DataType, typename IndexType = index_t>
CK_TILE_HOST void reference_grouped_topk(const HostTensor<DataType>& x,
HostTensor<DataType>& y_values,
HostTensor<IndexType>& 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<DataType> group_scores({num_token, num_expert_group}, {num_expert_group, 1});
HostTensor<DataType> group_mask({num_token, num_expert_group}, {num_expert_group, 1});
HostTensor<DataType> score_mask({num_token, num_expert}, {num_expert, 1});
HostTensor<DataType> masked_scores({num_token, num_expert}, {num_expert, 1});
HostTensor<DataType> group_values({num_token, topk_group}, {topk_group, 1});
HostTensor<IndexType> 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<DataType>::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<DataType, IndexType>(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<DataType>::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<DataType, IndexType>(masked_scores, y_values, y_indices, topk, dim, largest, sorted);
}
} // namespace ck_tile