cuda: fuse MoE weighted expert reduction (#25952)
* cuda : fuse MoE weighted reduction (mul + view + add) The MoE combine tail currently writes weighted expert outputs to global memory before reducing them. That intermediate global-memory traffic is the main cost. The production baseline generally runs two physical fused kernels; this path runs one. This change matches the full expert-weighting plus ordered-reduction subgraph and replaces it with one weighted-reduction kernel. Supported graphs: - unscaled: experts * router_weights - scaled: (experts * expert_scale) * router_weights k = 2..15 is handled by one runtime-k kernel. Matching is structural: op sequence, shapes, strides, expert views, and the left-to-right ADD chain. The fused kernel keeps that same reduction order. Results are not claimed bit-identical; CUDA FP32 contraction can change rounding slightly. Allocator integration uses add_alloc_dep from the graph-optimizer API so experts, router weights, and optional expert scales stay live until the fused destination is written. Memory ranges are rechecked before the fused kernel runs. Unrecognized or unsafe graphs are left alone and keep the existing per-op path. Set GGML_CUDA_MOE_WEIGHTED_REDUCTION=0 to disable the fusion. test-backend-ops covers scaled/unscaled, aligned/unaligned, and representative values across k=2..15, plus a k=16 case that must stay on the per-op path. * Pruned the test matrix from 15 to 6 * Addressed the aman and olivers review comments
This commit is contained in:
@@ -6474,6 +6474,79 @@ struct test_topk_moe : public test_case {
|
||||
}
|
||||
};
|
||||
|
||||
struct test_moe_weighted_reduction : public test_case {
|
||||
const int64_t n_embd;
|
||||
const int64_t n_expert_used;
|
||||
const int64_t n_tokens;
|
||||
const bool unaligned_experts;
|
||||
const bool with_expert_scale;
|
||||
const bool interleaved_views_adds;
|
||||
|
||||
test_moe_weighted_reduction(
|
||||
int64_t n_embd, int64_t n_expert_used, int64_t n_tokens,
|
||||
bool unaligned_experts = false, bool with_expert_scale = false, bool interleaved_views_adds = false) :
|
||||
n_embd(n_embd), n_expert_used(n_expert_used), n_tokens(n_tokens),
|
||||
unaligned_experts(unaligned_experts), with_expert_scale(with_expert_scale),
|
||||
interleaved_views_adds(interleaved_views_adds) {}
|
||||
|
||||
std::string vars() override {
|
||||
return VARS_TO_STR6(n_embd, n_expert_used, n_tokens, unaligned_experts, with_expert_scale, interleaved_views_adds);
|
||||
}
|
||||
|
||||
std::string op_desc(ggml_tensor * t) override {
|
||||
GGML_UNUSED(t);
|
||||
return "MOE_WEIGHTED_REDUCTION";
|
||||
}
|
||||
|
||||
bool run_whole_graph() override { return true; }
|
||||
|
||||
ggml_tensor * build_graph(ggml_context * ctx) override {
|
||||
ggml_tensor * experts;
|
||||
if (unaligned_experts) {
|
||||
ggml_tensor * storage = ggml_new_tensor_1d(
|
||||
ctx, GGML_TYPE_F32, n_embd * n_expert_used * n_tokens + 1);
|
||||
ggml_set_name(storage, "experts_storage");
|
||||
experts = ggml_view_3d(ctx, storage, n_embd, n_expert_used, n_tokens,
|
||||
n_embd * sizeof(float), n_embd * n_expert_used * sizeof(float), sizeof(float));
|
||||
} else {
|
||||
experts = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, n_expert_used, n_tokens);
|
||||
}
|
||||
ggml_set_name(experts, "experts");
|
||||
ggml_tensor * weights = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, n_expert_used, n_tokens);
|
||||
ggml_set_name(weights, "weights");
|
||||
|
||||
ggml_tensor * scaled = experts;
|
||||
if (with_expert_scale) {
|
||||
ggml_tensor * expert_scale = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, n_expert_used, n_tokens);
|
||||
ggml_set_name(expert_scale, "expert_scale");
|
||||
scaled = ggml_mul(ctx, experts, expert_scale);
|
||||
ggml_set_name(scaled, "scaled_experts");
|
||||
}
|
||||
|
||||
ggml_tensor * weighted = ggml_mul(ctx, scaled, weights);
|
||||
ggml_set_name(weighted, "weighted_experts");
|
||||
|
||||
std::vector<ggml_tensor *> views(n_expert_used);
|
||||
for (int64_t expert = 0; expert < n_expert_used; ++expert) {
|
||||
views[expert] = ggml_view_2d(
|
||||
ctx, weighted, n_embd, n_tokens, weighted->nb[2], expert * weighted->nb[1]);
|
||||
if (!interleaved_views_adds && mode == MODE_TEST) {
|
||||
ggml_build_forward_expand(gf, views[expert]);
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor * out = views[0];
|
||||
for (int64_t expert = 1; expert < n_expert_used; ++expert) {
|
||||
out = ggml_add(ctx, out, views[expert]);
|
||||
if (!interleaved_views_adds && mode == MODE_TEST) {
|
||||
ggml_build_forward_expand(gf, out);
|
||||
}
|
||||
}
|
||||
ggml_set_name(out, "moe_weighted_reduction");
|
||||
return out;
|
||||
}
|
||||
};
|
||||
|
||||
struct test_mul_mat_vec_fusion : public test_case {
|
||||
const ggml_type type;
|
||||
const ggml_glu_op glu_op;
|
||||
@@ -10268,6 +10341,14 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
}
|
||||
}
|
||||
|
||||
// Cover the supported boundaries, common k = 8 shapes, interleaved views and adds, and k = 16 fallback.
|
||||
test_cases.emplace_back(new test_moe_weighted_reduction(63, 2, 17));
|
||||
test_cases.emplace_back(new test_moe_weighted_reduction(2048, 8, 128));
|
||||
test_cases.emplace_back(new test_moe_weighted_reduction(2048, 8, 128, false, true));
|
||||
test_cases.emplace_back(new test_moe_weighted_reduction(63, 12, 33, true, true, true));
|
||||
test_cases.emplace_back(new test_moe_weighted_reduction(2048, 15, 40, false, true));
|
||||
test_cases.emplace_back(new test_moe_weighted_reduction(2048, 16, 32, false, true));
|
||||
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 1, 1));
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 16, 1, 1));
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 16, 1, 1, 1, true, true));
|
||||
|
||||
Reference in New Issue
Block a user