From 1dcb19084e4199e92f5c54a6a348fe2acda477a3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Sun, 20 Sep 2026 16:06:17 +0200 Subject: [PATCH] fuse grouped experts --- ggml/src/ggml-cuda/ggml-cuda.cu | 351 ++++++++++++++++++++++++++++++-- ggml/src/ggml-cuda/topk-moe.cu | 297 +++++++++++++++++++++++---- ggml/src/ggml-cuda/topk-moe.cuh | 5 + 3 files changed, 594 insertions(+), 59 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index a9038f1f4d..0935262598 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -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(), diff --git a/ggml/src/ggml-cuda/topk-moe.cu b/ggml/src/ggml-cuda/topk-moe.cu index dadcd601cb..68aa374bea 100644 --- a/ggml/src/ggml-cuda/topk-moe.cu +++ b/ggml/src/ggml-cuda/topk-moe.cu @@ -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 +__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( + 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(ctx, logits_d, weights_d, ids_d, bias_d, n_rows, n_experts, n_expert_used, clamp_val, diff --git a/ggml/src/ggml-cuda/topk-moe.cuh b/ggml/src/ggml-cuda/topk-moe.cuh index 061b37e297..fa51e3aa42 100644 --- a/ggml/src/ggml-cuda/topk-moe.cuh +++ b/ggml/src/ggml-cuda/topk-moe.cuh @@ -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,