add tests

This commit is contained in:
Sigbjørn Skjæret
2026-09-20 16:47:58 +02:00
committed by GitHub
parent 1dcb19084e
commit 0d0304f149
+32 -3
View File
@@ -6748,6 +6748,7 @@ struct test_topk_moe : public test_case {
const bool bias_probs;
const MoeGatingFunc gating_func;
const float scale_w;
const int n_expert_groups;
ggml_tensor * weights {};
ggml_tensor * selected_experts {};
@@ -6756,17 +6757,21 @@ struct test_topk_moe : public test_case {
bool with_norm = false,
bool bias_probs = false,
MoeGatingFunc gating_func = GATING_FUNC_SOFTMAX,
float scale_w = 0.0f) :
float scale_w = 0.0f,
int n_expert_groups = 1) :
ne(ne),
n_expert_used(n_expert_used),
with_norm(with_norm),
bias_probs(bias_probs),
gating_func(gating_func),
scale_w(scale_w) {
scale_w(scale_w),
n_expert_groups(n_expert_groups) {
GGML_ASSERT(n_expert_used <= ne[0]);
GGML_ASSERT(ne[0] % n_expert_groups == 0);
GGML_ASSERT(n_expert_used <= ne[0] / n_expert_groups);
}
std::string vars() override { return VARS_TO_STR6(ne, n_expert_used, with_norm, bias_probs, gating_func, scale_w); }
std::string vars() override { return VARS_TO_STR7(ne, n_expert_used, with_norm, bias_probs, gating_func, scale_w, n_expert_groups); }
std::string op_desc(ggml_tensor * t) override {
GGML_UNUSED(t);
@@ -6778,6 +6783,8 @@ struct test_topk_moe : public test_case {
ggml_tensor * build_graph(ggml_context * ctx) override {
const int n_expert = ne[0];
const int n_tokens = ne[1];
const int n_exp_per_group = n_expert / n_expert_groups;
const int n_group_used = n_expert_groups / 2;
ggml_tensor * logits = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data());
ggml_tensor * probs =
@@ -6794,6 +6801,24 @@ struct test_topk_moe : public test_case {
ggml_set_name(selection_probs, "selection_probs");
}
if (n_expert_groups > 1) {
ggml_tensor * selection_groups = ggml_reshape_3d(ctx, selection_probs, n_exp_per_group, n_expert_groups, n_tokens); // [n_exp_per_group, n_expert_groups, n_tokens]
ggml_tensor * group_scores = ggml_argsort_top_k(ctx, selection_groups, 2); // [2, n_expert_groups, n_tokens]
group_scores = ggml_get_rows(ctx, ggml_reshape_4d(ctx, selection_groups, 1, selection_groups->ne[0], selection_groups->ne[1], selection_groups->ne[2]), group_scores); // [1, 2, n_expert_groups, n_tokens]
group_scores = ggml_sum_rows(ctx, ggml_reshape_3d(ctx, group_scores, group_scores->ne[1], group_scores->ne[2], group_scores->ne[3])); // [1, n_expert_groups, n_tokens]
group_scores = ggml_reshape_2d(ctx, group_scores, group_scores->ne[1], group_scores->ne[2]); // [n_expert_groups, n_tokens]
ggml_tensor * expert_groups = ggml_argsort_top_k(ctx, group_scores, n_group_used); // [n_group_used, n_tokens]
ggml_set_name(expert_groups, "ffn_moe_group_topk");
selection_probs = ggml_get_rows(ctx, selection_groups, expert_groups); // [n_exp_per_group, n_group_used, n_tokens]
selection_probs = ggml_set_rows(ctx, ggml_fill(ctx, selection_groups, -INFINITY), selection_probs, expert_groups); // [n_exp_per_group, n_expert_groups, n_tokens]
selection_probs = ggml_reshape_2d(ctx, selection_probs, n_expert, n_tokens); // [n_expert, n_tokens]
ggml_set_name(selection_probs, "ffn_moe_probs_masked");
}
selected_experts = ggml_argsort_top_k(ctx, selection_probs, n_expert_used); // [n_expert_used, n_tokens]
ggml_set_name(selected_experts, "selected_experts");
@@ -10949,6 +10974,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_topk_moe({32, 8, 1, 1}, 4, with_norm, bias_probs, gate, scale_w));
test_cases.emplace_back(new test_topk_moe({32, 8, 1, 1}, 8, with_norm, bias_probs, gate, scale_w));
test_cases.emplace_back(new test_topk_moe({32, 9, 1, 1}, 8, with_norm, bias_probs, gate, scale_w));
// grouped experts
test_cases.emplace_back(new test_topk_moe({128, 1, 1, 1}, 8, with_norm, bias_probs, gate, scale_w, 16));
test_cases.emplace_back(new test_topk_moe({256, 22, 1, 1}, 6, with_norm, bias_probs, gate, scale_w, 32));
}
}
}