upstream-archive/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-25 15:37:25 -05:00
Code Issues Packages Projects Releases Wiki Activity

fuse grouped experts

Browse Source
This commit is contained in:
Sigbjørn Skjæret
2026-09-20 16:06:17 +02:00
committed by GitHub
parent 3d82ef62d4
commit 1dcb19084e
3 changed files with 594 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/ggml-cuda.cu
+336 -15
View File
@@ -2813,6 +2813,261 @@ static int ggml_cuda_try_gdn_cache_fusion(
return skip;
}
static bool ggml_cuda_topk_moe_grouped_experts(
const struct ggml_cgraph * cgraph,
int & node_idx,
const ggml_tensor * selection_src,
const ggml_tensor * probs_reshaped,
ggml_cuda_topk_moe_args & args) {
const int n_nodes = cgraph->n_nodes;
ggml_tensor ** nodes = cgraph->nodes;
// group reshape: [n_exp_per_group, n_expert_groups, n_tokens]
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_RESHAPE ||
nodes[node_idx]->src[0] != selection_src) {
return false;
}
const ggml_tensor * gr = nodes[node_idx];
const int64_t E = probs_reshaped->ne[1];
const int64_t P = gr->ne[0];
const int64_t G = gr->ne[1];
const int64_t T = gr->ne[2];
if (probs_reshaped->ne[0] != 1 ||
probs_reshaped->ne[2] != T ||
probs_reshaped->ne[3] != 1) {
return false;
}
if (P <= 0 || G <= 1 || T <= 0 || E <= 0 || P * G != E) {
return false;
}
// The current kernel assigns one thread per group.
if (G > WARP_SIZE) {
return false;
}
++node_idx;
// reshape to [1, n_exp_per_group, n_expert_groups, n_tokens]
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_RESHAPE ||
nodes[node_idx]->src[0] != gr) {
return false;
}
const ggml_tensor * r4 = nodes[node_idx];
if (r4->ne[0] != 1 || r4->ne[1] != P || r4->ne[2] != G || r4->ne[3] != T) {
return false;
}
++node_idx;
// argsort over experts within groups
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_ARGSORT ||
nodes[node_idx]->src[0] != gr) {
return false;
}
const ggml_tensor * ga_arg = nodes[node_idx];
++node_idx;
// view top-2 per group
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_VIEW ||
nodes[node_idx]->src[0] != ga_arg) {
return false;
}
const ggml_tensor * g2 = nodes[node_idx];
if (g2->ne[0] != 2 ||
g2->ne[1] != G || g2->ne[2] != T) {
return false;
}
++node_idx;
// gather top-k expert values per group
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_GET_ROWS) {
return false;
}
const ggml_tensor * get1 = nodes[node_idx];
if (get1->src[0] != r4 || get1->src[1] != g2) {
return false;
}
if (get1->ne[0] != 1 ||
get1->ne[1] != 2 ||
get1->ne[2] != G ||
get1->ne[3] != T) {
return false;
}
++node_idx;
// reshape to [2, n_expert_groups, n_tokens]
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_RESHAPE ||
nodes[node_idx]->src[0] != get1) {
return false;
}
const ggml_tensor * r5 = nodes[node_idx];
if (r5->ne[0] != 2 || r5->ne[1] != G || r5->ne[2] != T) {
return false;
}
++node_idx;
// sum top-k expert values per group
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_SUM_ROWS ||
nodes[node_idx]->src[0] != r5) {
return false;
}
const ggml_tensor * sum = nodes[node_idx];
if (sum->ne[0] != 1 || sum->ne[1] != G || sum->ne[2] != T) {
return false;
}
++node_idx;
// reshape group scores to [n_expert_groups, n_tokens]
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_RESHAPE ||
nodes[node_idx]->src[0] != sum) {
return false;
}
const ggml_tensor * r6 = nodes[node_idx];
if (r6->ne[0] != G || r6->ne[1] != T || r6->ne[2] != 1) {
return false;
}
++node_idx;
// argsort groups
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_ARGSORT ||
nodes[node_idx]->src[0] != r6) {
return false;
}
const ggml_tensor * earg = nodes[node_idx];
++node_idx;
// view top n_group_used groups
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_VIEW ||
nodes[node_idx]->src[0] != earg) {
return false;
}
const ggml_tensor * eg = nodes[node_idx];
if (eg->ne[0] < 1 || eg->ne[0] > G || eg->ne[1] != T) {
return false;
}
const int n_group_used = (int) eg->ne[0];
++node_idx;
// gather selected expert-group rows
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_GET_ROWS) {
return false;
}
const ggml_tensor * get2 = nodes[node_idx];
if (get2->src[0] != gr || get2->src[1] != eg) {
return false;
}
if (get2->ne[0] != P ||
get2->ne[1] != n_group_used ||
get2->ne[2] != T) {
return false;
}
++node_idx;
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_FILL ||
nodes[node_idx]->src[0] != gr) {
return false;
}
const ggml_tensor * fill = nodes[node_idx];
if (fill->ne[0] != P || fill->ne[1] != G || fill->ne[2] != T) {
return false;
}
if (ggml_get_op_params_f32(fill, 0) != -INFINITY) {
return false;
}
++node_idx;
// keep selected groups, -INFINITY elsewhere
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_SET_ROWS) {
return false;
}
const ggml_tensor * sel = nodes[node_idx];
if (sel->src[0] != get2 ||
sel->src[1] != eg ||
sel->src[2] != fill) {
return false;
}
if (sel->ne[0] != P || sel->ne[1] != G || sel->ne[2] != T) {
return false;
}
++node_idx;
// reshape masked scores back to [n_expert, n_tokens]
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_RESHAPE ||
nodes[node_idx]->src[0] != sel) {
return false;
}
const ggml_tensor * r7 = nodes[node_idx];
if (r7->ne[0] != E || r7->ne[1] != T || r7->ne[2] != 1) {
return false;
}
++node_idx;
args.grouped_experts = true;
args.n_expert_groups = (int) G;
args.n_exp_per_group = (int) P;
args.n_group_used = n_group_used;
return true;
}
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;
@@ -2820,6 +3075,7 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod
args.delayed_softmax = false;
args.prob_bias = false;
args.norm = false;
args.grouped_experts = false;
const int n_nodes = cgraph->n_nodes;
ggml_tensor ** nodes = cgraph->nodes;
@@ -2866,15 +3122,41 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod
args.prob_bias = true;
node_idx++;
}
// RESHAPE/ADD -> ARGSORT
if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_ARGSORT) {
return false;
const ggml_tensor * selection_src = args.prob_bias ? nodes[node_idx - 1] : nodes[node_idx - 2];
bool grouped_experts = false;
if (node_idx < n_nodes &&
nodes[node_idx]->op == GGML_OP_RESHAPE &&
nodes[node_idx]->src[0] == selection_src) {
const int saved_node_idx = node_idx;
grouped_experts = ggml_cuda_topk_moe_grouped_experts(cgraph, node_idx, selection_src, probs_reshaped, args);
if (!grouped_experts) {
node_idx = saved_node_idx;
}
}
if (args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
return false;
} else if (!args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 2]) {
return false;
if (!grouped_experts) {
// RESHAPE/ADD -> ARGSORT
if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_ARGSORT) {
return false;
}
if (args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
return false;
} else if (!args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 2]) {
return false;
}
} else {
// grouped path leaves node_idx at the final expert ARGSORT
if (node_idx >= n_nodes ||
nodes[node_idx]->op != GGML_OP_ARGSORT ||
nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
return false;
}
}
node_idx++;
@@ -3491,16 +3773,55 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
}
const int i_probs = i + (int) ops.size() - 1; // last node of the gating activation
if (args.prob_bias) {
bias = cgraph->nodes[i_probs + 2]->src[1];
ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW,
GGML_OP_GET_ROWS });
out_nodes[0] = i_probs + 4;
if (args.grouped_experts) {
// weights reshape: [1, n_expert, n_tokens]
ops.push_back(GGML_OP_RESHAPE);
if (args.prob_bias) {
ops.push_back(GGML_OP_ADD);
bias = cgraph->nodes[i + (int) ops.size() - 1]->src[1];
}
// grouped top-k sequence
ops.insert(ops.end(), {
GGML_OP_RESHAPE, // selection_groups
GGML_OP_RESHAPE, // 4d rows for get_rows
GGML_OP_ARGSORT, // top-2 per group
GGML_OP_VIEW, // top-2 per group view
GGML_OP_GET_ROWS, // gather top-2 per group values
GGML_OP_RESHAPE, // [2, n_groups, n_tokens]
GGML_OP_SUM_ROWS, // group score
GGML_OP_RESHAPE, // [n_groups, n_tokens]
GGML_OP_ARGSORT, // top groups
GGML_OP_VIEW, // top groups view
GGML_OP_GET_ROWS, // gather selected group rows
GGML_OP_FILL, // -INFINITY template
GGML_OP_SET_ROWS, // mask unselected groups
GGML_OP_RESHAPE // [n_expert, n_tokens]
});
// final expert top-k
ops.push_back(GGML_OP_ARGSORT);
const int view_idx = i + (int) ops.size();
ops.push_back(GGML_OP_VIEW);
ops.push_back(GGML_OP_GET_ROWS);
out_nodes[0] = view_idx;
ids = cgraph->nodes[out_nodes[0]];
} else {
ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS });
out_nodes[0] = i_probs + 3;
if (args.prob_bias) {
bias = cgraph->nodes[i_probs + 2]->src[1];
ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW,
GGML_OP_GET_ROWS });
out_nodes[0] = i_probs + 4;
} else {
ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS });
out_nodes[0] = i_probs + 3;
}
ids = cgraph->nodes[out_nodes[0]];
}
ids = cgraph->nodes[out_nodes[0]];
if (args.norm) {
ops.insert(ops.end(),
ggml/src/ggml-cuda/topk-moe.cu
+253 -44
View File
@@ -11,6 +11,11 @@ struct topk_moe_config {
bool use_sqrt_softplus;
bool with_norm;
bool delayed_softmax;
bool grouped_experts;
int n_expert_groups;
int n_exp_per_group;
int n_group_used;
};
// Warp-local softmax used for both the pre-top-k logits and the post-top-k delayed path.
@@ -78,12 +83,74 @@ __device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], co
}
}
__device__ __forceinline__ void topk_moe_top2_merge(float &a1, float &a2, float b1, float b2) {
// Merge two sorted top-2 pairs:
// a1 >= a2
// b1 >= b2
if (b1 > a1) {
a2 = max(a1, b2);
a1 = b1;
} else if (b1 > a2) {
a2 = b1;
}
if (b2 > a2) {
a2 = b2;
}
}
template <int experts_per_thread, bool has_bias>
__device__ __forceinline__ float topk_moe_group_score(const float (&wt)[experts_per_thread],
const float * bias,
const int n_experts,
const int n_exp_per_group,
const int group,
const int lane) {
const int start = group * n_exp_per_group;
const int end = start + n_exp_per_group;
float top1 = -INFINITY;
float top2 = -INFINITY;
#pragma unroll
for (int i = 0; i < experts_per_thread; ++i) {
const int e = i * WARP_SIZE + lane;
if (e >= start && e < end && e < n_experts) {
float v = wt[i];
if constexpr (has_bias) {
v += bias[e];
}
if (v > top1) {
top2 = top1;
top1 = v;
} else if (v > top2) {
top2 = v;
}
}
}
// Reduce the top-2 pair across the whole warp.
#pragma unroll
for (int mask = WARP_SIZE / 2; mask > 0; mask /= 2) {
const float o1 = __shfl_xor_sync(0xFFFFFFFF, top1, mask);
const float o2 = __shfl_xor_sync(0xFFFFFFFF, top2, mask);
topk_moe_top2_merge(top1, top2, o1, o2);
}
return (top2 == -INFINITY) ? top1 : (top1 + top2);
}
/*
This kernel does the following:
1. optionally softmax over the logits per token [n_experts, n_tokens]
2. argmax reduce over the top-k (n_experts_used) logits
3. write weights + ids to global memory
4. optionally normalize the weights or apply softmax over the selected logits
5. optionally select top-k grouped experts (n_group_used)
It is intended as fusion of softmax->top-k->get_rows pipeline for MoE models
*/
@@ -179,72 +246,208 @@ __global__ void topk_moe_cuda(const float * logits,
output_weights[i] = 0.f;
}
ggml_cuda_pdl_lc();
for (int k = 0; k < n_expert_used; k++) {
float max_val = wt[0];
int max_expert = threadIdx.x;
if (config.grouped_experts) {
const int G = config.n_expert_groups;
const int P = config.n_exp_per_group;
if constexpr (has_bias) {
float max_val_s = selection_wt[0];
// The score for each group must be reduced across the whole warp because
// a group is distributed over many lanes.
float gval = -INFINITY;
for (int g = 0; g < G; ++g) {
const float score = topk_moe_group_score<experts_per_thread, has_bias>(
wt, bias, n_experts, P, g, threadIdx.x
);
if (threadIdx.x == g) {
gval = score;
}
}
unsigned selected_mask = 0;
// Select top n_group_used groups.
for (int k = 0; k < config.n_group_used; ++k) {
float best = -INFINITY;
int best_group = WARP_SIZE;
int my_group = WARP_SIZE;
float my_score = -INFINITY;
if (threadIdx.x < G) {
// Skip groups that were already selected in a previous round.
if ((selected_mask & (1u << threadIdx.x)) == 0) {
my_group = threadIdx.x;
my_score = gval;
}
}
if (my_score > best || (my_score == best && my_group < best_group)) {
best = my_score;
best_group = my_group;
}
#pragma unroll
for (int i = 1; i < experts_per_thread; i++) {
for (int mask = WARP_SIZE / 2; mask > 0; mask /= 2) {
const float other = __shfl_xor_sync(0xFFFFFFFF, best, mask);
const int other_group = __shfl_xor_sync(0xFFFFFFFF, best_group, mask);
if (other > best || (other == best && other_group < best_group)) {
best = other;
best_group = other_group;
}
}
if (best_group < G) {
selected_mask |= (1u << best_group);
if (threadIdx.x == best_group) {
gval = -INFINITY;
}
}
}
// Build the masked expert selection scores.
// If bias is present, selection uses wt + bias, but output weights use wt.
float sel[experts_per_thread];
#pragma unroll
for (int i = 0; i < experts_per_thread; ++i) {
const int e = i * WARP_SIZE + threadIdx.x;
if ((n_experts % WARP_SIZE == 0 || e < n_experts) && P > 0) {
const int group = e / P;
if (group < G && ((selected_mask >> group) & 1u)) {
float v = wt[i];
if constexpr (has_bias) {
v += bias[e];
}
sel[i] = v;
} else {
sel[i] = -INFINITY;
}
} else {
sel[i] = -INFINITY;
}
}
ggml_cuda_pdl_lc();
// Final top-k expert selection over masked scores.
for (int k = 0; k < n_expert_used; ++k) {
float max_w = -INFINITY;
float max_s = -INFINITY;
int max_expert = threadIdx.x;
#pragma unroll
for (int i = 0; i < experts_per_thread; ++i) {
const int expert = threadIdx.x + i * WARP_SIZE;
if ((n_experts % WARP_SIZE == 0 || expert < n_experts) && selection_wt[i] > max_val_s) {
max_val = wt[i];
max_val_s = selection_wt[i];
if ((n_experts % WARP_SIZE == 0 || expert < n_experts) && sel[i] > max_s) {
max_s = sel[i];
max_w = wt[i];
max_expert = expert;
}
}
#pragma unroll
for (int mask = WARP_SIZE / 2; mask > 0; mask /= 2) {
const float val = __shfl_xor_sync(0xFFFFFFFF, max_val, mask, WARP_SIZE);
const float val_s = __shfl_xor_sync(0xFFFFFFFF, max_val_s, mask, WARP_SIZE);
const int expert = __shfl_xor_sync(0xFFFFFFFF, max_expert, mask, WARP_SIZE);
if (val_s > max_val_s || (val_s == max_val_s && expert < max_expert)) {
max_val = val;
max_val_s = val_s;
max_expert = expert;
const float other_s = __shfl_xor_sync(0xFFFFFFFF, max_s, mask);
const float other_w = __shfl_xor_sync(0xFFFFFFFF, max_w, mask);
const int other_exp = __shfl_xor_sync(0xFFFFFFFF, max_expert, mask);
if (other_s > max_s || (other_s == max_s && other_exp < max_expert)) {
max_s = other_s;
max_w = other_w;
max_expert = other_exp;
}
}
if ((k & (WARP_SIZE - 1)) == threadIdx.x) {
output_weights[k / WARP_SIZE] = max_w;
}
if ((max_expert & (WARP_SIZE - 1)) == threadIdx.x) {
selection_wt[max_expert / WARP_SIZE] = -INFINITY;
ids[k] = max_expert;
if (config.with_norm) {
wt_sum += max_w;
}
sel[max_expert / WARP_SIZE] = -INFINITY;
}
} else {
}
} else {
ggml_cuda_pdl_lc();
for (int k = 0; k < n_expert_used; k++) {
float max_val = wt[0];
int max_expert = threadIdx.x;
if constexpr (has_bias) {
float max_val_s = selection_wt[0];
#pragma unroll
for (int i = 1; i < experts_per_thread; i++) {
const int expert = threadIdx.x + i * WARP_SIZE;
if ((n_experts % WARP_SIZE == 0 || expert < n_experts) && wt[i] > max_val) {
max_val = wt[i];
max_expert = expert;
for (int i = 1; i < experts_per_thread; i++) {
const int expert = threadIdx.x + i * WARP_SIZE;
if ((n_experts % WARP_SIZE == 0 || expert < n_experts) && selection_wt[i] > max_val_s) {
max_val = wt[i];
max_val_s = selection_wt[i];
max_expert = expert;
}
}
#pragma unroll
for (int mask = WARP_SIZE / 2; mask > 0; mask /= 2) {
const float val = __shfl_xor_sync(0xFFFFFFFF, max_val, mask, WARP_SIZE);
const float val_s = __shfl_xor_sync(0xFFFFFFFF, max_val_s, mask, WARP_SIZE);
const int expert = __shfl_xor_sync(0xFFFFFFFF, max_expert, mask, WARP_SIZE);
if (val_s > max_val_s || (val_s == max_val_s && expert < max_expert)) {
max_val = val;
max_val_s = val_s;
max_expert = expert;
}
}
if ((max_expert & (WARP_SIZE - 1)) == threadIdx.x) {
selection_wt[max_expert / WARP_SIZE] = -INFINITY;
}
} else {
#pragma unroll
for (int i = 1; i < experts_per_thread; i++) {
const int expert = threadIdx.x + i * WARP_SIZE;
if ((n_experts % WARP_SIZE == 0 || expert < n_experts) && wt[i] > max_val) {
max_val = wt[i];
max_expert = expert;
}
}
#pragma unroll
for (int mask = WARP_SIZE / 2; mask > 0; mask /= 2) {
const float val = __shfl_xor_sync(0xFFFFFFFF, max_val, mask, WARP_SIZE);
const int expert = __shfl_xor_sync(0xFFFFFFFF, max_expert, mask, WARP_SIZE);
if (val > max_val || (val == max_val && expert < max_expert)) {
max_val = val;
max_expert = expert;
}
}
if ((max_expert & (WARP_SIZE - 1)) == threadIdx.x) {
wt[max_expert / WARP_SIZE] = -INFINITY;
}
}
#pragma unroll
for (int mask = WARP_SIZE / 2; mask > 0; mask /= 2) {
const float val = __shfl_xor_sync(0xFFFFFFFF, max_val, mask, WARP_SIZE);
const int expert = __shfl_xor_sync(0xFFFFFFFF, max_expert, mask, WARP_SIZE);
if (val > max_val || (val == max_val && expert < max_expert)) {
max_val = val;
max_expert = expert;
}
if ((k & (WARP_SIZE - 1)) == threadIdx.x) {
output_weights[k / WARP_SIZE] = max_val;
}
if ((max_expert & (WARP_SIZE - 1)) == threadIdx.x) {
wt[max_expert / WARP_SIZE] = -INFINITY;
}
}
if ((k & (WARP_SIZE - 1)) == threadIdx.x) {
output_weights[k / WARP_SIZE] = max_val;
}
if ((max_expert & (WARP_SIZE - 1)) == threadIdx.x) {
ids[k] = max_expert;
if (config.with_norm) {
wt_sum += max_val;
ids[k] = max_expert;
if (config.with_norm) {
wt_sum += max_val;
}
}
}
}
@@ -385,6 +588,12 @@ void ggml_cuda_op_topk_moe(ggml_backend_cuda_context & ctx,
config.use_sqrt_softplus = args.sqrt_softplus;
config.with_norm = with_norm;
config.delayed_softmax = args.delayed_softmax;
config.grouped_experts = args.grouped_experts;
config.n_expert_groups = args.n_expert_groups;
config.n_exp_per_group = args.n_exp_per_group;
config.n_group_used = args.n_group_used;
GGML_ASSERT(!(config.grouped_experts && config.delayed_softmax));
if (bias) {
launch_topk_moe_cuda<true>(ctx, logits_d, weights_d, ids_d, bias_d, n_rows, n_experts, n_expert_used, clamp_val,
ggml/src/ggml-cuda/topk-moe.cuh
+5
View File
@@ -14,6 +14,11 @@ struct ggml_cuda_topk_moe_args {
bool prob_bias{};
bool norm{};
bool scale{};
bool grouped_experts{};
int n_expert_groups{};
int n_exp_per_group{};
int n_group_used{};
};
void ggml_cuda_op_topk_moe(ggml_backend_cuda_context & ctx,