Skip to content

Commit a85c1e0

Browse files
committed
test(paged): cover MoE weighted combine chain
Assisted-by: Claude:opus-4.8 [Claude Code] Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
1 parent fd920cf commit a85c1e0

1 file changed

Lines changed: 90 additions & 0 deletions

File tree

tests/test-backend-ops.cpp

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4532,6 +4532,89 @@ struct test_moe_swiglu_down : public test_case {
45324532
}
45334533
};
45344534

4535+
// MoE down projection -> router-weight multiply -> rank-ordered expert add.
4536+
struct test_moe_weighted_combine : public test_case {
4537+
const ggml_type type_a;
4538+
const int n_mats;
4539+
const int n_used;
4540+
const int64_t n_ff;
4541+
const int64_t n_tokens;
4542+
const int64_t n_embd;
4543+
4544+
std::string vars() override {
4545+
return VARS_TO_STR6(type_a, n_mats, n_used, n_ff, n_tokens, n_embd);
4546+
}
4547+
4548+
double max_nmse_err() override {
4549+
return 5e-4;
4550+
}
4551+
4552+
double max_nmse_err(ggml_backend_t backend) override {
4553+
if ((type_a == GGML_TYPE_MXFP4 || type_a == GGML_TYPE_NVFP4) && backend_has_feature(backend, "BLACKWELL_NATIVE_FP4")) {
4554+
return 2e-2;
4555+
}
4556+
return max_nmse_err();
4557+
}
4558+
4559+
uint64_t op_flops(ggml_tensor * t) override {
4560+
GGML_UNUSED(t);
4561+
return 2 * n_ff * n_embd * n_tokens * n_used + 2 * n_embd * n_tokens * n_used;
4562+
}
4563+
4564+
test_moe_weighted_combine(ggml_type type_a = GGML_TYPE_F32, int n_mats = 128, int n_used = 8,
4565+
int64_t n_ff = 768, int64_t n_tokens = 128, int64_t n_embd = 2048)
4566+
: type_a(type_a), n_mats(n_mats), n_used(n_used), n_ff(n_ff), n_tokens(n_tokens), n_embd(n_embd) {
4567+
GGML_ASSERT(n_used <= n_mats);
4568+
}
4569+
4570+
ggml_tensor * build_graph(ggml_context * ctx) override {
4571+
ggml_tensor * down = ggml_new_tensor_3d(ctx, type_a, n_ff, n_embd, n_mats);
4572+
ggml_set_name(down, "down");
4573+
4574+
ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_mats, n_tokens);
4575+
ggml_set_name(ids, "ids");
4576+
if (n_used != n_mats) {
4577+
ids = ggml_view_2d(ctx, ids, n_used, n_tokens, ids->nb[1], 0);
4578+
ggml_set_name(ids, "view_of_ids");
4579+
}
4580+
4581+
ggml_tensor * act = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_ff, n_used, n_tokens);
4582+
ggml_set_name(act, "act");
4583+
4584+
ggml_tensor * weights = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, n_used, n_tokens);
4585+
ggml_set_name(weights, "weights");
4586+
4587+
ggml_tensor * experts = ggml_mul_mat_id(ctx, down, act, ids);
4588+
ggml_set_name(experts, "down_out");
4589+
4590+
experts = ggml_mul(ctx, experts, weights);
4591+
ggml_set_name(experts, "weighted");
4592+
4593+
ggml_tensor * out = ggml_view_2d(ctx, experts, n_embd, n_tokens, experts->nb[2], 0);
4594+
ggml_set_name(out, "rank_0");
4595+
4596+
for (int i = 1; i < n_used; ++i) {
4597+
ggml_tensor * rank = ggml_view_2d(ctx, experts, n_embd, n_tokens, experts->nb[2], i*experts->nb[1]);
4598+
ggml_set_name(rank, "rank_i");
4599+
out = ggml_add(ctx, out, rank);
4600+
ggml_set_name(out, "rank_sum");
4601+
}
4602+
4603+
return out;
4604+
}
4605+
4606+
void initialize_tensors(ggml_context * ctx) override {
4607+
init_mul_mat_id_tensors(ctx, n_mats);
4608+
}
4609+
4610+
bool run_whole_graph() override { return true; }
4611+
4612+
std::string op_desc(ggml_tensor * t) override {
4613+
GGML_UNUSED(t);
4614+
return "MOE_WEIGHTED_COMBINE";
4615+
}
4616+
};
4617+
45354618
// GGML_OP_OUT_PROD
45364619
struct test_out_prod : public test_case {
45374620
const ggml_type type_a;
@@ -8851,6 +8934,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
88518934
test_cases.emplace_back(new test_moe_swiglu_down(GGML_TYPE_NVFP4, 128, 8, 768, n, 2048));
88528935
}
88538936

8937+
// [paged Phase 7] MoE down projection -> router-weight multiply -> rank-ordered
8938+
// expert add gate for the weighted-combine fusion candidate.
8939+
test_cases.emplace_back(new test_moe_weighted_combine(GGML_TYPE_F32, 8, 2, 32, 8, 64));
8940+
for (int n : {16, 33, 64, 128, 130, 200}) {
8941+
test_cases.emplace_back(new test_moe_weighted_combine(GGML_TYPE_NVFP4, 128, 8, 768, n, 2048));
8942+
}
8943+
88548944
// [paged P0 / track B] NVFP4/MXFP4 dense decode-shape mmq_y-down bit-exact gate.
88558945
// The dense FP4 weight GEMM is the track-B target; P1 lowers mmq_y (the weight-row tile) on the
88568946
// NVFP4 decode path to raise resident-CTA occupancy. mmq_y is a pure N-row tiling knob, so a

0 commit comments

Comments
 (0)