Skip to content

Commit 54ba263

Browse files
committed
test-backend-ops : make experts more evenly probable (test_moe)
1 parent b0b83dd commit 54ba263

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

tests/test-backend-ops.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1172,7 +1172,7 @@ struct test_moe : public test_case {
11721172
ggml_tensor * cur = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens);
11731173

11741174
ggml_tensor * logits = ggml_mul_mat(ctx, ffn_gate_inp, cur); // [n_tokens, num_experts]
1175-
ggml_tensor * probs = ggml_soft_max(ctx, logits); // [n_tokens, num_experts]
1175+
ggml_tensor * probs = ggml_soft_max_ext(ctx, logits, nullptr, 1.0f/sqrtf(n_embd)); // [n_tokens, num_experts]
11761176

11771177
// select experts
11781178
ggml_tensor * selected_experts = ggml_top_k(ctx, probs, n_experts_per_tok); // [n_tokens, num_experts_per_tok]

0 commit comments

Comments
 (0)