Superminaren/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-04 02:37:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

cuda: fuse MoE weighted expert reduction (#25952)

Browse Source 
* 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:
anujj
2026-09-01 21:48:47 +02:00
committed by GitHub
parent b356fa2624
commit 3466812d1f
4 changed files with 332 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/ggml-cuda.cu
+179 -2
View File
@@ -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"
@@ -3026,6 +3027,150 @@ 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) {
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 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);
};
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]) && is_weights(mul->src[0])) {
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<ggml_op> ops(node_count, GGML_OP_VIEW);
ops[0] = GGML_OP_MUL;
if (mul_count == 2) {
ops[1] = GGML_OP_MUL;
}
std::vector<const ggml_tensor *> 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;
}
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_can_fuse(const struct ggml_cgraph * cgraph,
int node_idx,
@@ -3288,6 +3433,18 @@ 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) {
ggml_cuda_moe_weighted_reduction_match match;
if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) {
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;
}
}
}
// 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;
@@ -4340,10 +4497,30 @@ 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;
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)) {
continue;
}
params->add_alloc_dep(params->user_data, const_cast<ggml_tensor *>(match.experts), match.dst);
params->add_alloc_dep(params->user_data, const_cast<ggml_tensor *>(match.weights), match.dst);
if (match.expert_scale != nullptr) {
params->add_alloc_dep(
params->user_data, const_cast<ggml_tensor *>(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);
ggml/src/ggml-cuda/moe-weighted-reduction.cu
+65
View File
@@ -0,0 +1,65 @@
#include "moe-weighted-reduction.cuh"
static __global__ void moe_weighted_reduction_f32(const float * __restrict__ experts,
const float * __restrict__ expert_scale,
const float * __restrict__ weights,
float * __restrict__ dst,
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 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 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[token * n_embd + col] = 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 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_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());
}
ggml/src/ggml-cuda/moe-weighted-reduction.cuh
+7
View File
@@ -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);