From eb24eb8aed33bfeb0282a68bafb1274bcee8a217 Mon Sep 17 00:00:00 2001 From: Anuj Jalota Date: Mon, 31 Aug 2026 23:09:38 +0530 Subject: [PATCH 1/3] 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. --- ggml/src/ggml-cuda/ggml-cuda.cu | 209 +++++++++++++++++- ggml/src/ggml-cuda/moe-weighted-reduction.cu | 72 ++++++ ggml/src/ggml-cuda/moe-weighted-reduction.cuh | 7 + tests/test-backend-ops.cpp | 84 +++++++ 4 files changed, 367 insertions(+), 5 deletions(-) create mode 100644 ggml/src/ggml-cuda/moe-weighted-reduction.cu create mode 100644 ggml/src/ggml-cuda/moe-weighted-reduction.cuh diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index bd9754c2ffdf..fdbac5a5fcaa 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -32,6 +32,7 @@ #include "ggml-cuda/mmq.cuh" #include "ggml-cuda/mmvf.cuh" #include "ggml-cuda/mmvq.cuh" +#include "ggml-cuda/moe-weighted-reduction.cuh" #include "ggml-cuda/norm.cuh" #include "ggml-cuda/opt-step-adamw.cuh" #include "ggml-cuda/opt-step-sgd.cuh" @@ -2957,7 +2958,7 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod return true; } -// returns whether the write (out) nodes overwrite the read nodes in operation +// Returns whether the write (out) nodes overwrite the read nodes in operation. static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, const int node_idx, const int node_count, @@ -3021,6 +3022,177 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, return is_ok; } +// The long form spans 2*k + 1 nodes. ggml_can_fuse_subgraph() accepts at most +// 31 nodes, so k <= 15; larger values use the per-operation path. +static constexpr int MOE_WEIGHTED_REDUCTION_MAX_EXPERTS = 15; + +struct ggml_cuda_moe_weighted_reduction_match { + const ggml_tensor * experts = nullptr; + const ggml_tensor * expert_scale = nullptr; + const ggml_tensor * weights = nullptr; + ggml_tensor * dst = nullptr; + int node_count = 0; +}; + +static bool ggml_cuda_match_moe_weighted_reduction( + const ggml_cgraph * cgraph, + int node_idx, + ggml_cuda_moe_weighted_reduction_match & match, + bool check_memory_ranges = true) { + const ggml_tensor * first = cgraph->nodes[node_idx]; + if (first->op != GGML_OP_MUL || first->type != GGML_TYPE_F32 || !ggml_is_contiguous(first)) { + return false; + } + + auto is_weights = [](const ggml_tensor * tensor, const ggml_tensor * full) { + return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && + tensor->ne[0] == 1 && tensor->ne[1] == full->ne[1] && + tensor->ne[2] == full->ne[2] && tensor->ne[3] == full->ne[3]; + }; + auto is_experts = [](const ggml_tensor * tensor, const ggml_tensor * full) { + return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && + ggml_are_same_shape(tensor, full); + }; + + auto split_mul = [&is_experts, &is_weights](const ggml_tensor * mul, + const ggml_tensor *& full, + const ggml_tensor *& broadcast) { + if (is_experts(mul->src[0], mul) && is_weights(mul->src[1], mul)) { + full = mul->src[0]; + broadcast = mul->src[1]; + return true; + } + if (is_experts(mul->src[1], mul) && is_weights(mul->src[0], mul)) { + full = mul->src[1]; + broadcast = mul->src[0]; + return true; + } + return false; + }; + + const ggml_tensor * weighted = first; + const ggml_tensor * experts = nullptr; + const ggml_tensor * expert_scale = nullptr; + const ggml_tensor * weights = nullptr; + int mul_count = 1; + + // Match both structural forms: + // (experts * expert_scale) * router_weight + // experts * router_weight + // The matcher does not depend on the model or quantization type. + if (node_idx + 1 < cgraph->n_nodes) { + const ggml_tensor * second = cgraph->nodes[node_idx + 1]; + const ggml_tensor * scaled = nullptr; + const ggml_tensor * route = nullptr; + const ggml_tensor * raw = nullptr; + const ggml_tensor * scale = nullptr; + if (second->op == GGML_OP_MUL && second->type == GGML_TYPE_F32 && ggml_is_contiguous(second) && + split_mul(second, scaled, route) && scaled == first && split_mul(first, raw, scale)) { + weighted = second; + experts = raw; + expert_scale = scale; + weights = route; + mul_count = 2; + } + } + + if (experts == nullptr && !split_mul(first, experts, weights)) { + return false; + } + + const int n_expert_used = (int) weighted->ne[1]; + const int64_t n_tokens = weighted->ne[2] * weighted->ne[3]; + if (n_expert_used < 2 || n_expert_used > MOE_WEIGHTED_REDUCTION_MAX_EXPERTS || n_tokens <= 0) { + return false; + } + + const int node_count = 2 * n_expert_used + mul_count - 1; + if (node_idx + node_count > cgraph->n_nodes) { + return false; + } + + std::vector ops(node_count, GGML_OP_VIEW); + ops[0] = GGML_OP_MUL; + if (mul_count == 2) { + ops[1] = GGML_OP_MUL; + } + std::vector views; + views.reserve(n_expert_used); + const ggml_tensor * previous = nullptr; + int n_adds = 0; + for (int offset = mul_count; offset < node_count; ++offset) { + const ggml_tensor * candidate = cgraph->nodes[node_idx + offset]; + ops[offset] = candidate->op; + + if (candidate->op == GGML_OP_VIEW) { + const int expert = (int) views.size(); + if (expert >= n_expert_used || candidate->src[0] != weighted || candidate->view_src != weighted || + candidate->type != GGML_TYPE_F32 || candidate->ne[0] != weighted->ne[0] || + candidate->ne[1] != n_tokens || candidate->ne[2] != 1 || candidate->ne[3] != 1 || + candidate->nb[0] != weighted->nb[0] || candidate->nb[1] != weighted->nb[2] || + candidate->view_offs != (size_t) expert * weighted->nb[1]) { + return false; + } + views.push_back(candidate); + continue; + } + + if (candidate->op != GGML_OP_ADD || views.size() < 2 || n_adds + 1 >= (int) views.size()) { + return false; + } + const ggml_tensor * lhs = n_adds == 0 ? views[0] : previous; + const ggml_tensor * rhs = views[n_adds + 1]; + if (candidate->src[0] != lhs || candidate->src[1] != rhs || candidate->type != GGML_TYPE_F32) { + return false; + } + previous = candidate; + ++n_adds; + } + + if ((int) views.size() != n_expert_used || n_adds != n_expert_used - 1 || previous == nullptr) { + return false; + } + if (!ggml_is_contiguous(previous) || previous->ne[0] != weighted->ne[0] || + previous->ne[1] != n_tokens || previous->ne[2] != 1 || previous->ne[3] != 1) { + return false; + } + + const int output_idx = node_idx + node_count - 1; + if (!ggml_can_fuse_subgraph(cgraph, node_idx, node_count, ops.data(), &output_idx, 1)) { + return false; + } + + if (check_memory_ranges && + !ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, node_count, &output_idx, 1)) { + return false; + } + + match.experts = experts; + match.expert_scale = expert_scale; + match.weights = weights; + match.dst = cgraph->nodes[output_idx]; + match.node_count = node_count; + return true; +} + +static bool ggml_cuda_use_moe_weighted_reduction() { + static const bool enabled = [] { + // Enabled by default. Unrecognized graphs use the per-operation path. + // Set GGML_CUDA_MOE_WEIGHTED_REDUCTION=0 to disable the fusion. + const char * env = getenv("GGML_CUDA_MOE_WEIGHTED_REDUCTION"); + return env == nullptr || atoi(env) != 0; + }(); + return enabled; +} + +static bool ggml_cuda_fusion_disabled() { + static const bool disabled = [] { + const char * env = getenv("GGML_CUDA_DISABLE_FUSION"); + return env != nullptr && atoi(env) != 0; + }(); + return disabled; +} + static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, int node_idx, @@ -3276,13 +3448,21 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, // try and fuse nodes and return the number of nodes to skip static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { - static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); - if (disable_fusion) { + if (ggml_cuda_fusion_disabled()) { return 0; } ggml_tensor * node = cgraph->nodes[i]; + if (ggml_cuda_use_moe_weighted_reduction() && node->op == GGML_OP_MUL) { + ggml_cuda_moe_weighted_reduction_match match; + if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { + ggml_cuda_op_moe_weighted_reduction( + *cuda_ctx, match.experts, match.expert_scale, match.weights, match.dst); + return match.node_count - 1; + } + } + // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache if (node->op == GGML_OP_GATED_DELTA_NET) { ggml_cuda_gated_delta_net_fused_cache fused_state_cpy; @@ -4329,10 +4509,29 @@ static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_ev } static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) { - GGML_UNUSED(params); - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; + if (!ggml_cuda_fusion_disabled() && ggml_cuda_use_moe_weighted_reduction()) { + for (int i = 0; i < cgraph->n_nodes; ++i) { + if (cgraph->nodes[i]->op != GGML_OP_MUL) { + continue; + } + + ggml_cuda_moe_weighted_reduction_match match; + if (!ggml_cuda_match_moe_weighted_reduction(cgraph, i, match, false)) { + continue; + } + + params->add_alloc_dep(params->user_data, const_cast(match.experts), match.dst); + params->add_alloc_dep(params->user_data, const_cast(match.weights), match.dst); + if (match.expert_scale != nullptr) { + params->add_alloc_dep( + params->user_data, const_cast(match.expert_scale), match.dst); + } + i += match.node_count - 1; + } + } + #ifdef USE_CUDA_GRAPH const void * graph_key = ggml_cuda_graph_get_key(cgraph); const bool use_cuda_graph = ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); diff --git a/ggml/src/ggml-cuda/moe-weighted-reduction.cu b/ggml/src/ggml-cuda/moe-weighted-reduction.cu new file mode 100644 index 000000000000..1327dce0feb6 --- /dev/null +++ b/ggml/src/ggml-cuda/moe-weighted-reduction.cu @@ -0,0 +1,72 @@ +#include "moe-weighted-reduction.cuh" + +#include + +static __global__ void moe_weighted_reduction_f32(const float * __restrict__ experts, + const float * __restrict__ expert_scale, + const float * __restrict__ weights, + float * __restrict__ dst, + int64_t n_embd, + int64_t n_tokens, + int n_expert_used) { + const int64_t index = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + const int64_t total = n_embd * n_tokens; + if (index >= total) { + return; + } + + const int64_t token = index / n_embd; + const int64_t col = index - token * n_embd; + const int64_t first_row = token * n_expert_used; + const float first_scale = expert_scale != nullptr ? expert_scale[first_row] : 1.0f; + float sum = (experts[first_row * n_embd + col] * first_scale) * weights[first_row]; + + for (int expert = 1; expert < n_expert_used; ++expert) { + const int64_t row = token * n_expert_used + expert; + const float scale = expert_scale != nullptr ? expert_scale[row] : 1.0f; + sum += (experts[row * n_embd + col] * scale) * weights[row]; + } + dst[index] = sum; +} + +static void launch_moe_weighted_reduction(const float * experts, + const float * expert_scale, + const float * weights, + float * dst, + int64_t n_embd, + int64_t n_tokens, + int n_expert_used, + cudaStream_t stream) { + constexpr int threads = 256; + const int64_t blocks = (n_embd * n_tokens + threads - 1) / threads; + GGML_ASSERT(blocks <= INT_MAX); + moe_weighted_reduction_f32 + <<(blocks), threads, 0, stream>>>( + experts, expert_scale, weights, dst, n_embd, n_tokens, n_expert_used); +} + +void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx, + const ggml_tensor * experts, + const ggml_tensor * expert_scale, + const ggml_tensor * weights, + ggml_tensor * dst) { + GGML_ASSERT(experts->type == GGML_TYPE_F32); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(expert_scale == nullptr || expert_scale->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(experts)); + GGML_ASSERT(ggml_is_contiguous(weights)); + GGML_ASSERT(expert_scale == nullptr || ggml_is_contiguous(expert_scale)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int64_t n_embd = experts->ne[0]; + const int64_t n_expert_used = experts->ne[1]; + const int64_t n_tokens = experts->ne[2] * experts->ne[3]; + cudaStream_t stream = ctx.stream(); + + launch_moe_weighted_reduction((const float *) experts->data, + expert_scale ? (const float *) expert_scale->data : nullptr, + (const float *) weights->data, + (float *) dst->data, n_embd, n_tokens, (int) n_expert_used, stream); + CUDA_CHECK(cudaGetLastError()); +} diff --git a/ggml/src/ggml-cuda/moe-weighted-reduction.cuh b/ggml/src/ggml-cuda/moe-weighted-reduction.cuh new file mode 100644 index 000000000000..b72f947ab398 --- /dev/null +++ b/ggml/src/ggml-cuda/moe-weighted-reduction.cuh @@ -0,0 +1,7 @@ +#include "common.cuh" + +void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx, + const ggml_tensor * experts, + const ggml_tensor * expert_scale, + const ggml_tensor * weights, + ggml_tensor * dst); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 4a7a0623174c..e388819c8ec6 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -6336,6 +6336,71 @@ 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; + + 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) : + n_embd(n_embd), n_expert_used(n_expert_used), n_tokens(n_tokens), + unaligned_experts(unaligned_experts), with_expert_scale(with_expert_scale) {} + + std::string vars() override { + return VARS_TO_STR5(n_embd, n_expert_used, n_tokens, unaligned_experts, with_expert_scale); + } + + 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 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]); + } + + ggml_tensor * out = views[0]; + for (int64_t expert = 1; expert < n_expert_used; ++expert) { + out = ggml_add(ctx, out, views[expert]); + } + 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; @@ -10075,6 +10140,25 @@ static std::vector> make_test_cases_eval() { } } + // Exercise k=2..8 with aligned/unaligned and scaled/unscaled inputs. + test_cases.emplace_back(new test_moe_weighted_reduction(63, 2, 17)); + test_cases.emplace_back(new test_moe_weighted_reduction(256, 3, 33)); + test_cases.emplace_back(new test_moe_weighted_reduction(256, 4, 33)); + test_cases.emplace_back(new test_moe_weighted_reduction(256, 4, 33, true)); + test_cases.emplace_back(new test_moe_weighted_reduction(63, 5, 17, false, true)); + test_cases.emplace_back(new test_moe_weighted_reduction(256, 6, 33, true)); + test_cases.emplace_back(new test_moe_weighted_reduction(256, 7, 33, false, true)); + 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, 4, 17, true, true)); + // Exercise k=9..15 with aligned/unaligned and scaled/unscaled inputs. + test_cases.emplace_back(new test_moe_weighted_reduction(2048, 9, 65)); // aligned expert storage, unscaled + test_cases.emplace_back(new test_moe_weighted_reduction(2048, 15, 40, false, true)); // aligned expert storage, scaled (max fusable k) + test_cases.emplace_back(new test_moe_weighted_reduction(256, 9, 48, true, false)); // unaligned expert storage, unscaled + test_cases.emplace_back(new test_moe_weighted_reduction(63, 12, 33, true, true)); // non-multiple-of-four embedding size, scaled + // k > 15 exceeds the 31-node fusion limit and must use the per-operation path. + test_cases.emplace_back(new test_moe_weighted_reduction(2048, 16, 32, false, true)); // not fused; correctness via the per-operation path + 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)); From 727432158177ec346e1a37cf7b65e482a7f63746 Mon Sep 17 00:00:00 2001 From: Anuj Jalota Date: Tue, 1 Sep 2026 11:48:46 +0530 Subject: [PATCH 2/3] Pruned the test matrix from 15 to 6 --- tests/test-backend-ops.cpp | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index e388819c8ec6..17d3c2a556e8 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10140,24 +10140,13 @@ static std::vector> make_test_cases_eval() { } } - // Exercise k=2..8 with aligned/unaligned and scaled/unscaled inputs. + // Cover the supported boundaries, common k = 8 shapes, a view-backed input, 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(256, 3, 33)); - test_cases.emplace_back(new test_moe_weighted_reduction(256, 4, 33)); - test_cases.emplace_back(new test_moe_weighted_reduction(256, 4, 33, true)); - test_cases.emplace_back(new test_moe_weighted_reduction(63, 5, 17, false, true)); - test_cases.emplace_back(new test_moe_weighted_reduction(256, 6, 33, true)); - test_cases.emplace_back(new test_moe_weighted_reduction(256, 7, 33, false, true)); 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, 4, 17, true, true)); - // Exercise k=9..15 with aligned/unaligned and scaled/unscaled inputs. - test_cases.emplace_back(new test_moe_weighted_reduction(2048, 9, 65)); // aligned expert storage, unscaled - test_cases.emplace_back(new test_moe_weighted_reduction(2048, 15, 40, false, true)); // aligned expert storage, scaled (max fusable k) - test_cases.emplace_back(new test_moe_weighted_reduction(256, 9, 48, true, false)); // unaligned expert storage, unscaled - test_cases.emplace_back(new test_moe_weighted_reduction(63, 12, 33, true, true)); // non-multiple-of-four embedding size, scaled - // k > 15 exceeds the 31-node fusion limit and must use the per-operation path. - test_cases.emplace_back(new test_moe_weighted_reduction(2048, 16, 32, false, true)); // not fused; correctness via the per-operation path + test_cases.emplace_back(new test_moe_weighted_reduction(63, 12, 33, 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)); From 7a6227435934b28678a4963a813df3d81003d1b0 Mon Sep 17 00:00:00 2001 From: Anuj Jalota Date: Tue, 1 Sep 2026 22:08:32 +0530 Subject: [PATCH 3/3] Addressed the aman and olivers review comments --- ggml/src/ggml-cuda/ggml-cuda.cu | 76 +++++++------------- ggml/src/ggml-cuda/moe-weighted-reduction.cu | 31 ++++---- tests/test-backend-ops.cpp | 18 +++-- 3 files changed, 52 insertions(+), 73 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index fdbac5a5fcaa..a773a4845597 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2958,7 +2958,7 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod return true; } -// Returns whether the write (out) nodes overwrite the read nodes in operation. +// returns whether the write (out) nodes overwrite the read nodes in operation static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, const int node_idx, const int node_count, @@ -3037,32 +3037,28 @@ struct ggml_cuda_moe_weighted_reduction_match { static bool ggml_cuda_match_moe_weighted_reduction( const ggml_cgraph * cgraph, int node_idx, - ggml_cuda_moe_weighted_reduction_match & match, - bool check_memory_ranges = true) { + ggml_cuda_moe_weighted_reduction_match & match) { const ggml_tensor * first = cgraph->nodes[node_idx]; if (first->op != GGML_OP_MUL || first->type != GGML_TYPE_F32 || !ggml_is_contiguous(first)) { return false; } - auto is_weights = [](const ggml_tensor * tensor, const ggml_tensor * full) { - return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && - tensor->ne[0] == 1 && tensor->ne[1] == full->ne[1] && - tensor->ne[2] == full->ne[2] && tensor->ne[3] == full->ne[3]; - }; - auto is_experts = [](const ggml_tensor * tensor, const ggml_tensor * full) { - return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && - ggml_are_same_shape(tensor, full); - }; + auto split_mul = [](const ggml_tensor * mul, const ggml_tensor *& full, const ggml_tensor *& broadcast) { + auto is_weights = [mul](const ggml_tensor * tensor) { + return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && tensor->ne[0] == 1 && + tensor->ne[1] == mul->ne[1] && tensor->ne[2] == mul->ne[2] && tensor->ne[3] == mul->ne[3]; + }; + auto is_experts = [mul](const ggml_tensor * tensor) { + return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && + ggml_are_same_shape(tensor, mul); + }; - auto split_mul = [&is_experts, &is_weights](const ggml_tensor * mul, - const ggml_tensor *& full, - const ggml_tensor *& broadcast) { - if (is_experts(mul->src[0], mul) && is_weights(mul->src[1], mul)) { + if (is_experts(mul->src[0]) && is_weights(mul->src[1])) { full = mul->src[0]; broadcast = mul->src[1]; return true; } - if (is_experts(mul->src[1], mul) && is_weights(mul->src[0], mul)) { + if (is_experts(mul->src[1]) && is_weights(mul->src[0])) { full = mul->src[1]; broadcast = mul->src[0]; return true; @@ -3100,8 +3096,8 @@ static bool ggml_cuda_match_moe_weighted_reduction( return false; } - const int n_expert_used = (int) weighted->ne[1]; - const int64_t n_tokens = weighted->ne[2] * weighted->ne[3]; + const int n_expert_used = (int) weighted->ne[1]; + const int64_t n_tokens = weighted->ne[2] * weighted->ne[3]; if (n_expert_used < 2 || n_expert_used > MOE_WEIGHTED_REDUCTION_MAX_EXPERTS || n_tokens <= 0) { return false; } @@ -3162,11 +3158,6 @@ static bool ggml_cuda_match_moe_weighted_reduction( return false; } - if (check_memory_ranges && - !ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, node_count, &output_idx, 1)) { - return false; - } - match.experts = experts; match.expert_scale = expert_scale; match.weights = weights; @@ -3175,24 +3166,6 @@ static bool ggml_cuda_match_moe_weighted_reduction( return true; } -static bool ggml_cuda_use_moe_weighted_reduction() { - static const bool enabled = [] { - // Enabled by default. Unrecognized graphs use the per-operation path. - // Set GGML_CUDA_MOE_WEIGHTED_REDUCTION=0 to disable the fusion. - const char * env = getenv("GGML_CUDA_MOE_WEIGHTED_REDUCTION"); - return env == nullptr || atoi(env) != 0; - }(); - return enabled; -} - -static bool ggml_cuda_fusion_disabled() { - static const bool disabled = [] { - const char * env = getenv("GGML_CUDA_DISABLE_FUSION"); - return env != nullptr && atoi(env) != 0; - }(); - return disabled; -} - static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, int node_idx, @@ -3448,18 +3421,22 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, // try and fuse nodes and return the number of nodes to skip static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { - if (ggml_cuda_fusion_disabled()) { + static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); + if (disable_fusion) { return 0; } ggml_tensor * node = cgraph->nodes[i]; - if (ggml_cuda_use_moe_weighted_reduction() && node->op == GGML_OP_MUL) { + if (node->op == GGML_OP_MUL) { ggml_cuda_moe_weighted_reduction_match match; if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { - ggml_cuda_op_moe_weighted_reduction( - *cuda_ctx, match.experts, match.expert_scale, match.weights, match.dst); - return match.node_count - 1; + const int output_idx = i + match.node_count - 1; + if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, match.node_count, &output_idx, 1)) { + ggml_cuda_op_moe_weighted_reduction( + *cuda_ctx, match.experts, match.expert_scale, match.weights, match.dst); + return match.node_count - 1; + } } } @@ -4511,14 +4488,15 @@ static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_ev static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) { ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; - if (!ggml_cuda_fusion_disabled() && ggml_cuda_use_moe_weighted_reduction()) { + static const bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); + if (!disable_fusion) { for (int i = 0; i < cgraph->n_nodes; ++i) { if (cgraph->nodes[i]->op != GGML_OP_MUL) { continue; } ggml_cuda_moe_weighted_reduction_match match; - if (!ggml_cuda_match_moe_weighted_reduction(cgraph, i, match, false)) { + if (!ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { continue; } diff --git a/ggml/src/ggml-cuda/moe-weighted-reduction.cu b/ggml/src/ggml-cuda/moe-weighted-reduction.cu index 1327dce0feb6..11ec58497f1e 100644 --- a/ggml/src/ggml-cuda/moe-weighted-reduction.cu +++ b/ggml/src/ggml-cuda/moe-weighted-reduction.cu @@ -1,32 +1,27 @@ #include "moe-weighted-reduction.cuh" -#include - static __global__ void moe_weighted_reduction_f32(const float * __restrict__ experts, const float * __restrict__ expert_scale, const float * __restrict__ weights, float * __restrict__ dst, - int64_t n_embd, - int64_t n_tokens, - int n_expert_used) { - const int64_t index = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; - const int64_t total = n_embd * n_tokens; - if (index >= total) { + const int64_t n_embd, + const int n_expert_used) { + const int64_t token = blockIdx.x; + const int64_t col = (int64_t) blockIdx.y * blockDim.x + threadIdx.x; + if (col >= n_embd) { return; } - const int64_t token = index / n_embd; - const int64_t col = index - token * n_embd; - const int64_t first_row = token * n_expert_used; - const float first_scale = expert_scale != nullptr ? expert_scale[first_row] : 1.0f; - float sum = (experts[first_row * n_embd + col] * first_scale) * weights[first_row]; + const uint64_t first_row = (uint64_t) token * n_expert_used; + const float first_scale = expert_scale != nullptr ? expert_scale[first_row] : 1.0f; + float sum = (experts[first_row * n_embd + col] * first_scale) * weights[first_row]; for (int expert = 1; expert < n_expert_used; ++expert) { - const int64_t row = token * n_expert_used + expert; + const uint64_t row = first_row + expert; const float scale = expert_scale != nullptr ? expert_scale[row] : 1.0f; sum += (experts[row * n_embd + col] * scale) * weights[row]; } - dst[index] = sum; + dst[token * n_embd + col] = sum; } static void launch_moe_weighted_reduction(const float * experts, @@ -38,11 +33,9 @@ static void launch_moe_weighted_reduction(const float * experts, int n_expert_used, cudaStream_t stream) { constexpr int threads = 256; - const int64_t blocks = (n_embd * n_tokens + threads - 1) / threads; - GGML_ASSERT(blocks <= INT_MAX); + const dim3 blocks(n_tokens, (n_embd + threads - 1) / threads, 1); moe_weighted_reduction_f32 - <<(blocks), threads, 0, stream>>>( - experts, expert_scale, weights, dst, n_embd, n_tokens, n_expert_used); + <<>>(experts, expert_scale, weights, dst, n_embd, n_expert_used); } void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx, diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 17d3c2a556e8..a83fdad88ebc 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -6342,15 +6342,17 @@ struct test_moe_weighted_reduction : public test_case { 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 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) {} + unaligned_experts(unaligned_experts), with_expert_scale(with_expert_scale), + interleaved_views_adds(interleaved_views_adds) {} std::string vars() override { - return VARS_TO_STR5(n_embd, n_expert_used, n_tokens, unaligned_experts, with_expert_scale); + 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 { @@ -6390,11 +6392,17 @@ struct test_moe_weighted_reduction : public test_case { 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; @@ -10140,11 +10148,11 @@ static std::vector> make_test_cases_eval() { } } - // Cover the supported boundaries, common k = 8 shapes, a view-backed input, and k = 16 fallback. + // 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)); + 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));