From acf9c9e3cd82d8d58ff4bdf731f11379a77adca2 Mon Sep 17 00:00:00 2001 From: Josh Leverette Date: Sun, 26 Jul 2026 20:37:58 -0500 Subject: [PATCH] cuda: fuse MoE expert weighting and reduction MoE output currently applies router weights with one or two MUL nodes, creates one view per selected expert, and reduces those views through an ADD chain. Recognize that subgraph and replace it with one CUDA kernel that applies the factors and accumulates experts in the original order. Write directly to the graph output when allocations do not overlap. When the allocator reuses an input range for the output, use a pooled temporary and one device copy so the fusion remains race-free. DGX Spark / GB10 pp8192 benchmark: llama-bench -p 8192 -n 0 -b 16384 -ub 2048 -r 7 -fa 1 -lm none Values are median prompt tokens/s. Each row compares upstream 0d47ea742 with this standalone change: north-mini-code NVFP4 3152.7 -> 3435.0 (+9.0%) qwen3.6-35b-a3b NVFP4 2863.9 -> 3070.3 (+7.2%) laguna-xs Q4_K 2856.1 -> 3016.7 (+5.6%) qwen3.6-35b-a3b BF16 1601.2 -> 1674.1 (+4.6%) qwen3.6-27b BF16 dense 938.9 -> 935.5 (-0.4%, noise) gemma-4-26b-a4b Q4_K 3097.9 -> 3287.4 (+6.1%) gpt-oss-20b MXFP4 4587.9 -> 4852.3 (+5.8%) gpt-oss-120b MXFP4 2403.6 -> 2491.9 (+3.7%) Every MoE format tested improves. Dense Qwen does not match the fusion and is performance-neutral. Deterministic output validation in the separately tested two-patch series showed that applying this reduction to the scheduler-only control was byte-identical for all eight models above. Laguna and Gemma used their Jinja chat templates. The standalone change builds successfully and passes all 73 NVFP4 MUL_MAT_ID backend tests. --- ggml/src/ggml-cuda/binbcast.cu | 130 ++++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/binbcast.cuh | 7 ++ ggml/src/ggml-cuda/ggml-cuda.cu | 87 +++++++++++++++++++++ 3 files changed, 224 insertions(+) diff --git a/ggml/src/ggml-cuda/binbcast.cu b/ggml/src/ggml-cuda/binbcast.cu index 2e38077bf67..010b94f38d4 100644 --- a/ggml/src/ggml-cuda/binbcast.cu +++ b/ggml/src/ggml-cuda/binbcast.cu @@ -542,6 +542,136 @@ void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst, } } +template +static __global__ void k_moe_weighted_sum( + const float * __restrict__ experts, + const float * __restrict__ scale, + const float * __restrict__ weights, + float * __restrict__ dst, + const int n_embd) { + ggml_cuda_pdl_lc(); + + const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int t = blockIdx.y; + + if (i >= n_embd) { + return; + } + + const size_t expert_base = size_t(t) * n_experts * n_embd + i; + const size_t factor_base = size_t(t) * n_experts; + + ggml_cuda_pdl_sync(); + + float value = experts[expert_base]; + if constexpr (has_scale) { + value = __fmul_rn(value, scale[factor_base]); + } + float sum = __fmul_rn(value, weights[factor_base]); + +#pragma unroll + for (int e = 1; e < n_experts; ++e) { + value = experts[expert_base + size_t(e) * n_embd]; + if constexpr (has_scale) { + value = __fmul_rn(value, scale[factor_base + e]); + } + value = __fmul_rn(value, weights[factor_base + e]); + sum = __fadd_rn(sum, value); + } + + dst[size_t(t) * n_embd + i] = sum; +} + +template +static void launch_moe_weighted_sum( + const float * experts, + const float * scale, + const float * weights, + float * dst, + const int n_embd, + const int n_tokens, + const int n_experts, + cudaStream_t stream) { + const dim3 block_dims(256, 1, 1); + const dim3 block_nums((n_embd + block_dims.x - 1) / block_dims.x, n_tokens, 1); + const ggml_cuda_kernel_launch_params launch_params(block_nums, block_dims, 0, stream); + +#define GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(N) \ + ggml_cuda_kernel_launch(k_moe_weighted_sum, launch_params, experts, scale, weights, dst, n_embd) + + switch (n_experts) { + case 2: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(2); break; + case 3: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(3); break; + case 4: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(4); break; + case 5: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(5); break; + case 6: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(6); break; + case 7: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(7); break; + case 8: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(8); break; + case 9: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(9); break; + case 10: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(10); break; + case 11: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(11); break; + case 12: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(12); break; + case 13: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(13); break; + case 14: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(14); break; + case 15: GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM(15); break; + default: GGML_ABORT("unsupported number of experts"); + } + +#undef GGML_CUDA_LAUNCH_MOE_WEIGHTED_SUM +} + +void ggml_cuda_op_moe_weighted_sum( + ggml_backend_cuda_context & ctx, + const ggml_tensor * experts, + const ggml_tensor * scale, + const ggml_tensor * weights, + ggml_tensor * dst) { + GGML_ASSERT(experts->type == GGML_TYPE_F32); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(scale == nullptr || scale->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(experts)); + GGML_ASSERT(ggml_is_contiguous(weights)); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(scale == nullptr || ggml_is_contiguous(scale)); + + const int n_embd = experts->ne[0]; + const int n_experts = experts->ne[1]; + const int n_tokens = experts->ne[2]; + + const auto overlaps = [](const ggml_tensor * a, const ggml_tensor * b) { + const uintptr_t a_begin = (uintptr_t) a->data; + const uintptr_t a_end = a_begin + ggml_nbytes(a); + const uintptr_t b_begin = (uintptr_t) b->data; + const uintptr_t b_end = b_begin + ggml_nbytes(b); + return a_begin < b_end && b_begin < a_end; + }; + + const bool needs_tmp = + overlaps(dst, experts) || overlaps(dst, weights) || (scale && overlaps(dst, scale)); + ggml_cuda_pool_alloc tmp(ctx.pool()); + float * dst_ptr = (float *) dst->data; + if (needs_tmp) { + dst_ptr = tmp.alloc(ggml_nelements(dst)); + } + + if (scale) { + launch_moe_weighted_sum( + (const float *) experts->data, (const float *) scale->data, (const float *) weights->data, + dst_ptr, + n_embd, n_tokens, n_experts, ctx.stream()); + } else { + launch_moe_weighted_sum( + (const float *) experts->data, nullptr, (const float *) weights->data, + dst_ptr, + n_embd, n_tokens, n_experts, ctx.stream()); + } + + if (needs_tmp) { + CUDA_CHECK(cudaMemcpyAsync(dst->data, dst_ptr, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, ctx.stream())); + } +} + void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; diff --git a/ggml/src/ggml-cuda/binbcast.cuh b/ggml/src/ggml-cuda/binbcast.cuh index 12624785b44..6e057c04aab 100644 --- a/ggml/src/ggml-cuda/binbcast.cuh +++ b/ggml/src/ggml-cuda/binbcast.cuh @@ -10,3 +10,10 @@ void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst void ggml_cuda_op_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse); void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse); + +void ggml_cuda_op_moe_weighted_sum( + ggml_backend_cuda_context & ctx, + const ggml_tensor * experts, + const ggml_tensor * scale, + const ggml_tensor * weights, + ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index e73a7b8906c..9e610e7eb30 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2701,6 +2701,86 @@ static int ggml_cuda_try_gdn_cache_fusion( return skip; } +static int ggml_cuda_try_moe_weighted_sum_fusion( + ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int node_idx) { + ggml_tensor * first_mul = cgraph->nodes[node_idx]; + if (first_mul->op != GGML_OP_MUL || first_mul->type != GGML_TYPE_F32 || + first_mul->ne[1] < 2 || first_mul->ne[1] > 15 || + first_mul->ne[2] < 1 || first_mul->ne[3] != 1) { + return 0; + } + + int n_muls = 1; + ggml_tensor * weighted = first_mul; + if (node_idx + 1 < cgraph->n_nodes && + cgraph->nodes[node_idx + 1]->op == GGML_OP_MUL && + cgraph->nodes[node_idx + 1]->src[0] == first_mul && + ggml_are_same_shape(cgraph->nodes[node_idx + 1], first_mul)) { + weighted = cgraph->nodes[node_idx + 1]; + n_muls = 2; + } + + const int n_experts = weighted->ne[1]; + const int view_idx = node_idx + n_muls; + const int add_idx = view_idx + n_experts; + const int n_nodes = n_muls + n_experts + n_experts - 1; + if (node_idx + n_nodes > cgraph->n_nodes) { + return 0; + } + + for (int e = 0; e < n_experts; ++e) { + const ggml_tensor * view = cgraph->nodes[view_idx + e]; + if (view->op != GGML_OP_VIEW || view->view_src != weighted || + view->view_offs != size_t(e) * weighted->nb[1] || + view->ne[0] != weighted->ne[0] || view->ne[1] != weighted->ne[2] || + view->ne[2] != 1 || view->ne[3] != 1) { + return 0; + } + } + + for (int e = 1; e < n_experts; ++e) { + const ggml_tensor * add = cgraph->nodes[add_idx + e - 1]; + const ggml_tensor * lhs = e == 1 ? cgraph->nodes[view_idx] : cgraph->nodes[add_idx + e - 2]; + const ggml_tensor * rhs = cgraph->nodes[view_idx + e]; + if (add->op != GGML_OP_ADD || add->src[0] != lhs || add->src[1] != rhs) { + return 0; + } + } + + const ggml_tensor * experts = first_mul->src[0]; + const ggml_tensor * scale = n_muls == 2 ? first_mul->src[1] : nullptr; + const ggml_tensor * weights = n_muls == 2 ? weighted->src[1] : first_mul->src[1]; + ggml_tensor * output = cgraph->nodes[node_idx + n_nodes - 1]; + + if (experts->type != GGML_TYPE_F32 || weights->type != GGML_TYPE_F32 || + (scale && scale->type != GGML_TYPE_F32) || + !ggml_is_contiguous(experts) || !ggml_is_contiguous(weights) || + (scale && !ggml_is_contiguous(scale)) || !ggml_is_contiguous(output) || + weights->ne[0] != 1 || weights->ne[1] != n_experts || weights->ne[2] != experts->ne[2] || + (scale && (scale->ne[0] != 1 || scale->ne[1] != n_experts || scale->ne[2] != experts->ne[2]))) { + return 0; + } + + std::vector ops(n_nodes); + for (int j = 0; j < n_muls; ++j) { + ops[j] = GGML_OP_MUL; + } + for (int j = 0; j < n_experts; ++j) { + ops[n_muls + j] = GGML_OP_VIEW; + } + for (int j = 0; j < n_experts - 1; ++j) { + ops[n_muls + n_experts + j] = GGML_OP_ADD; + } + + const int out_node = node_idx + n_nodes - 1; + if (!ggml_can_fuse_subgraph(cgraph, node_idx, n_nodes, ops.data(), &out_node, 1)) { + return 0; + } + + ggml_cuda_op_moe_weighted_sum(*cuda_ctx, experts, scale, weights, output); + return n_nodes - 1; +} + static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int node_idx, ggml_cuda_topk_moe_args & args) { args.sigmoid = false; args.sqrt_softplus = false; @@ -3150,6 +3230,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_tensor * node = cgraph->nodes[i]; + if (node->op == GGML_OP_MUL) { + const int nodes_to_skip = ggml_cuda_try_moe_weighted_sum_fusion(cuda_ctx, cgraph, i); + if (nodes_to_skip > 0) { + return nodes_to_skip; + } + } + // 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;