Commit ea40a3a1e for llama.cpp
commit ea40a3a1e1774534b7fe605caa1e85558175e3e8
Author: Georgi Gerganov <ggerganov@gmail.com>
Date: Thu Nov 5 13:52:56 2026 +0200
metal : fuse DeepSeek-V4/V4.1 and GLM MoE routing (#31615)
* metal : support scale-bias routing norm in the TOPK_MOE fusion
The DeepSeek-V4.1 graph normalizes the routing weights with a scale+bias
(sum + 1e-20) instead of the lower clamp used by the other models. Add
norm_bias/norm_scale_bias fusion chains that accept a SCALE node with a
unit scale factor at the clamp position, pass an additive bias in the
kernel args, and apply it as max(wt_sum, clamp) + bias so both variants
share one kernel. Cover the new chains in test-backend-ops.
Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (Mac Studio, M5 Ultra, 256GB)
* metal : fuse sqrt(softplus) MoE routing with a selection bias
Add TOPK_MOE chains anchored at the softplus node (UNARY+SQRT+ADD+
ARGSORT+GET_ROWS plus optional norm/scale) so the DeepSeek-V4/V4.1
routing fuses on Metal. The kernel scores with sqrt(softplus) and ranks
selection-biased keys while the gathered weights and the norm stay
unbiased, matching the reference graph. Dispatch hooks at the softplus
node, cover the norm-bias combination in test-backend-ops and record the
new deepseek4/deepseek41 fusion counters.
Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (Mac Studio, M5 Ultra, 256GB)
* metal : fuse sigmoid MoE routing with a selection bias
The sigmoid-gated MoE archs (GLM, DeepSeek-V2/V3.2, ...) rank
sigmoid(logits) + exp_probs_b with an optional weights-norm, but only
the softmax and sqrt(softplus) routing chains were fusable.
- add 6 sigmoid-anchored TOPK_MOE chains: sigmoid + ADD + ARGSORT +
GET_ROWS with optional SCALE and SUM_ROWS + CLAMP/SCALE + DIV norm
tail; the selection-bias ADD is required, same as the sqrt(softplus)
chains
- generalize the gating function constant from a sqrt(softplus) bool
to an int (0 = softmax, 1 = sqrt(softplus), 2 = sigmoid); sigmoid
scores are 1 / (1 + exp(-logits)), matching the CPU, Metal unary and
CUDA kernels
- anchor the unary dispatch hook at sigmoid nodes too
- test-backend-ops: 4 sigmoid routing cases with and without a
weights scale
- re-record tests/fusion/MTL.csv: the sigmoid archs (glm5-next,
deepseek2/32, kimi, minimax, ...) now fuse their routing, and
muse-spark prefill drops one MOE_REDUCE fire because the reordered
graph allocates the routing weights inside the reduce output range,
which the dst/input overlap check correctly declines
Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (Mac Studio, M5 Ultra, 256GB)
* cont : remove low-signal comments
* ocd : valign ggml_metal_fusions
diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp
index 463d57c7b..72133d86d 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-device.cpp
@@ -1738,12 +1738,13 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge(ggml
}
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_topk_moe(
- ggml_metal_library_t lib, int32_t n_expert, int32_t top_k, bool with_norm) {
+ ggml_metal_library_t lib, int32_t n_expert, int32_t top_k, bool with_norm, int32_t gating, bool pbias) {
char base[256];
char name[256];
snprintf(base, 256, "kernel_topk_moe_f32");
- snprintf(name, 256, "%s_n_expert=%d_top_k=%d_with_norm=%d", base, n_expert, top_k, with_norm ? 1 : 0);
+ snprintf(name, 256, "%s_n_expert=%d_top_k=%d_with_norm=%d_gating=%d_pbias=%d",
+ base, n_expert, top_k, with_norm ? 1 : 0, gating, pbias ? 1 : 0);
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
if (!res.pipeline) {
@@ -1751,6 +1752,8 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_topk_moe(
ggml_metal_cv_set_bool (cv, with_norm, FC_TOPK_MOE + 0);
ggml_metal_cv_set_int32(cv, n_expert, FC_TOPK_MOE + 1);
ggml_metal_cv_set_int32(cv, top_k, FC_TOPK_MOE + 2);
+ ggml_metal_cv_set_int32(cv, gating, FC_TOPK_MOE + 3);
+ ggml_metal_cv_set_bool (cv, pbias, FC_TOPK_MOE + 4);
res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h
index e096d611b..4053a624e 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.h
+++ b/ggml/src/ggml-metal/ggml-metal-device.h
@@ -154,7 +154,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge (ggml_metal_library_t lib, const struct ggml_tensor * op);
-struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_topk_moe (ggml_metal_library_t lib, int32_t n_expert, int32_t top_k, bool with_norm);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_topk_moe (ggml_metal_library_t lib, int32_t n_expert, int32_t top_k, bool with_norm, int32_t gating, bool pbias);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_moe_reduce (ggml_metal_library_t lib, int32_t n_expert_used);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t n_fuse);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin_one (ggml_metal_library_t lib, enum ggml_op op);
diff --git a/ggml/src/ggml-metal/ggml-metal-fusion.cpp b/ggml/src/ggml-metal/ggml-metal-fusion.cpp
index 020cad29d..6766b4bd9 100644
--- a/ggml/src/ggml-metal/ggml-metal-fusion.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-fusion.cpp
@@ -345,9 +345,6 @@ static bool ggml_metal_fusion_check_snake(
#define GGML_METAL_TOPK_MOE_MAX_EXPERTS 1024
-// SOFT_MAX + ARGSORT + GET_ROWS (plus optional norm/scale) for MoE routing.
-// This is a multi-output elision chain: the fused kernel writes both the selected
-// expert ids and the gathered/normalized routing weights.
static const std::vector<ggml_op> ops_topk_moe = {
GGML_OP_SOFT_MAX, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS
};
@@ -362,6 +359,60 @@ static const std::vector<ggml_op> ops_topk_moe_norm_scale = {
GGML_OP_SOFT_MAX, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS,
GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_RESHAPE, GGML_OP_SCALE
};
+static const std::vector<ggml_op> ops_topk_moe_norm_bias = {
+ GGML_OP_SOFT_MAX, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS,
+ GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_SCALE, GGML_OP_DIV, GGML_OP_RESHAPE
+};
+static const std::vector<ggml_op> ops_topk_moe_norm_scale_bias = {
+ GGML_OP_SOFT_MAX, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS,
+ GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_SCALE, GGML_OP_DIV, GGML_OP_RESHAPE, GGML_OP_SCALE
+};
+
+static const std::vector<ggml_op> ops_topk_moe_sp = {
+ GGML_OP_UNARY, GGML_OP_SQRT, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS
+};
+static const std::vector<ggml_op> ops_topk_moe_sp_scale = {
+ GGML_OP_UNARY, GGML_OP_SQRT, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS, GGML_OP_SCALE
+};
+static const std::vector<ggml_op> ops_topk_moe_sp_norm = {
+ GGML_OP_UNARY, GGML_OP_SQRT, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS,
+ GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_RESHAPE
+};
+static const std::vector<ggml_op> ops_topk_moe_sp_norm_scale = {
+ GGML_OP_UNARY, GGML_OP_SQRT, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS,
+ GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_RESHAPE, GGML_OP_SCALE
+};
+static const std::vector<ggml_op> ops_topk_moe_sp_norm_bias = {
+ GGML_OP_UNARY, GGML_OP_SQRT, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS,
+ GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_SCALE, GGML_OP_DIV, GGML_OP_RESHAPE
+};
+static const std::vector<ggml_op> ops_topk_moe_sp_norm_scale_bias = {
+ GGML_OP_UNARY, GGML_OP_SQRT, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS,
+ GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_SCALE, GGML_OP_DIV, GGML_OP_RESHAPE, GGML_OP_SCALE
+};
+
+static const std::vector<ggml_op> ops_topk_moe_sig = {
+ GGML_OP_UNARY, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS
+};
+static const std::vector<ggml_op> ops_topk_moe_sig_scale = {
+ GGML_OP_UNARY, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS, GGML_OP_SCALE
+};
+static const std::vector<ggml_op> ops_topk_moe_sig_norm = {
+ GGML_OP_UNARY, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS,
+ GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_RESHAPE
+};
+static const std::vector<ggml_op> ops_topk_moe_sig_norm_scale = {
+ GGML_OP_UNARY, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS,
+ GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_DIV, GGML_OP_RESHAPE, GGML_OP_SCALE
+};
+static const std::vector<ggml_op> ops_topk_moe_sig_norm_bias = {
+ GGML_OP_UNARY, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS,
+ GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_SCALE, GGML_OP_DIV, GGML_OP_RESHAPE
+};
+static const std::vector<ggml_op> ops_topk_moe_sig_norm_scale_bias = {
+ GGML_OP_UNARY, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS,
+ GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_SCALE, GGML_OP_DIV, GGML_OP_RESHAPE, GGML_OP_SCALE
+};
static bool ggml_metal_fusion_check_topk_moe(
const ggml_metal_fusion * fusion,
@@ -375,8 +426,11 @@ static bool ggml_metal_fusion_check_topk_moe(
const int n_ops = (int) fusion->ops.size();
- const bool with_norm = n_ops >= 6;
- const bool with_scale = n_ops == 4 || n_ops == 7;
+ bool with_norm = false;
+ for (int j = 0; j < n_ops; ++j) {
+ with_norm |= fusion->ops[j] == GGML_OP_SUM_ROWS;
+ }
+ const bool with_scale = fusion->ops[n_ops - 1] == GGML_OP_SCALE;
// the fusion table operates on the non-empty node sequence; the raw graph also
// contains the RESHAPE/VIEW nodes that the fused kernel elides.
@@ -411,24 +465,70 @@ static bool ggml_metal_fusion_check_topk_moe(
}
}
- const ggml_tensor * softmax = gf->nodes[raw_start];
- const ggml_tensor * probs_reshaped = gf->nodes[raw_start + 1];
- const ggml_tensor * argsort = gf->nodes[raw_start + 2];
- const ggml_tensor * ids = gf->nodes[raw_start + 3];
- const ggml_tensor * get_rows = gf->nodes[raw_start + 4];
- const ggml_tensor * out = gf->nodes[raw_end];
- const ggml_tensor * logits = softmax->src[0];
-
- // the fused kernel implements plain softmax only
- float scale = 1.0f;
- float max_bias = 0.0f;
- memcpy(&scale, ((const int32_t *) softmax->op_params) + 0, sizeof(scale));
- memcpy(&max_bias, ((const int32_t *) softmax->op_params) + 1, sizeof(max_bias));
- if (scale != 1.0f || max_bias != 0.0f || softmax->src[1] || softmax->src[2]) {
+ const ggml_tensor * first = gf->nodes[raw_start];
+ const ggml_tensor * out = gf->nodes[raw_end];
+
+ // the unary chains start at the sigmoid or the softplus node and rank biased scores
+ const bool with_unary = first->op == GGML_OP_UNARY;
+ const bool with_sp = with_unary && ggml_get_unary_op(first) == GGML_UNARY_OP_SOFTPLUS;
+ const bool with_sig = with_unary && ggml_get_unary_op(first) == GGML_UNARY_OP_SIGMOID;
+
+ if (with_unary && !with_sp && !with_sig) {
return false;
}
- if (logits->type != GGML_TYPE_F32 || softmax->type != GGML_TYPE_F32 ||
+ const ggml_tensor * probs; // the scoring output
+ const ggml_tensor * scores; // the node the argsort ranks (adds the bias for the unary chains)
+ const ggml_tensor * probs_reshaped;
+ const ggml_tensor * argsort;
+ const ggml_tensor * ids;
+ const ggml_tensor * get_rows;
+ const ggml_tensor * logits = first->src[0];
+ int ids_off;
+
+ if (with_unary) {
+ // the softplus chains add a SQRT node, the sigmoid chains start at the sigmoid node
+ const int base = with_sp ? 1 : 0;
+
+ probs = gf->nodes[raw_start + base + 0];
+ probs_reshaped = gf->nodes[raw_start + base + 1];
+ scores = gf->nodes[raw_start + base + 2];
+ argsort = gf->nodes[raw_start + base + 3];
+ ids = gf->nodes[raw_start + base + 4];
+ get_rows = gf->nodes[raw_start + base + 5];
+ ids_off = base + 4;
+
+ if (scores->op != GGML_OP_ADD || scores->src[0] != probs ||
+ (with_sp && (probs->op != GGML_OP_SQRT || probs->src[0] != first))) {
+ return false;
+ }
+
+ // the selection bias is a per-expert vector shared by all rows
+ const ggml_tensor * bias = scores->src[1];
+ if (bias->type != GGML_TYPE_F32 || scores->type != GGML_TYPE_F32 || !ggml_is_contiguous(bias) ||
+ bias->ne[0] != probs->ne[0] || bias->ne[1] != 1 || bias->ne[2] != 1 || bias->ne[3] != 1) {
+ return false;
+ }
+ } else {
+ probs = first;
+ scores = first;
+ probs_reshaped = gf->nodes[raw_start + 1];
+ argsort = gf->nodes[raw_start + 2];
+ ids = gf->nodes[raw_start + 3];
+ get_rows = gf->nodes[raw_start + 4];
+ ids_off = 3;
+
+ // the fused kernel implements plain softmax only
+ float scale = 1.0f;
+ float max_bias = 0.0f;
+ memcpy(&scale, ((const int32_t *) first->op_params) + 0, sizeof(scale));
+ memcpy(&max_bias, ((const int32_t *) first->op_params) + 1, sizeof(max_bias));
+ if (scale != 1.0f || max_bias != 0.0f || first->src[1] || first->src[2]) {
+ return false;
+ }
+ }
+
+ if (logits->type != GGML_TYPE_F32 || probs->type != GGML_TYPE_F32 ||
out->type != GGML_TYPE_F32 || ids->type != GGML_TYPE_I32) {
return false;
}
@@ -455,38 +555,40 @@ static bool ggml_metal_fusion_check_topk_moe(
return false;
}
- if (probs_reshaped->src[0] != softmax || argsort->src[0] != softmax ||
+ // the top-k gathers the unbiased scores, the argsort ranks the biased ones
+ if (probs_reshaped->src[0] != probs || argsort->src[0] != scores ||
ids->src[0] != argsort || get_rows->src[0] != probs_reshaped || get_rows->src[1] != ids) {
return false;
}
if (with_norm) {
- const ggml_tensor * weights_reshaped = gf->nodes[raw_start + 5];
- const ggml_tensor * sum_rows = gf->nodes[raw_start + 6];
- const ggml_tensor * clamp = gf->nodes[raw_start + 7];
- const ggml_tensor * div = gf->nodes[raw_start + 8];
- const ggml_tensor * out_reshaped = gf->nodes[raw_start + 9];
+ const int norm_off = with_unary ? 6 + (with_sp ? 1 : 0) : 5;
+
+ const ggml_tensor * weights_reshaped = gf->nodes[raw_start + norm_off + 0];
+ const ggml_tensor * sum_rows = gf->nodes[raw_start + norm_off + 1];
+ const ggml_tensor * denom = gf->nodes[raw_start + norm_off + 2];
+ const ggml_tensor * div = gf->nodes[raw_start + norm_off + 3];
+ const ggml_tensor * out_reshaped = gf->nodes[raw_start + norm_off + 4];
if (weights_reshaped->src[0] != get_rows || sum_rows->src[0] != weights_reshaped ||
- clamp->src[0] != sum_rows || div->src[0] != weights_reshaped || div->src[1] != clamp ||
+ denom->src[0] != sum_rows || div->src[0] != weights_reshaped || div->src[1] != denom ||
out_reshaped->src[0] != div) {
return false;
}
- if (with_scale) {
- const ggml_tensor * scale_node = gf->nodes[raw_start + 10];
- if (scale_node->src[0] != out_reshaped) {
- return false;
- }
+ // the kernel applies the sum adjustment as a lower clamp or an additive bias, never a scale
+ if (denom->op == GGML_OP_SCALE && ggml_get_op_params_f32(denom, 0) != 1.0f) {
+ return false;
}
- } else if (with_scale) {
- const ggml_tensor * scale_node = gf->nodes[raw_start + 5];
- if (scale_node->src[0] != get_rows) {
+
+ if (with_scale && out->src[0] != out_reshaped) {
return false;
}
+ } else if (with_scale && out->src[0] != get_rows) {
+ return false;
}
- const int outputs[2] = { raw_start + 3, raw_end };
+ const int outputs[2] = { raw_start + ids_off, raw_end };
if (!ggml_can_fuse_subgraph_ext(gf, raw_idxs, raw_count, ops_all.data(), outputs, 2)) {
return false;
}
@@ -763,33 +865,54 @@ static const std::vector<ggml_op> ops_moe_reduce_8 = {
static const std::vector<ggml_op> ops_mul_mat_add = { GGML_OP_MUL_MAT, GGML_OP_ADD };
static const std::vector<ggml_metal_fusion> ggml_metal_fusions = {
- { GGML_METAL_FUSION_NORM_MUL, ops_norm_mul, {}, false, ggml_metal_fusion_check_norm },
- { GGML_METAL_FUSION_NORM_MUL_ADD, ops_norm_mul_add, {}, false, ggml_metal_fusion_check_norm },
- { GGML_METAL_FUSION_NORM_SCALE, ops_norm_scale, {}, false, ggml_metal_fusion_check_norm },
- { GGML_METAL_FUSION_NORM_MUL, ops_rms_norm_mul, {}, false, ggml_metal_fusion_check_norm },
- { GGML_METAL_FUSION_NORM_MUL_ADD, ops_rms_norm_mul_add, {}, false, ggml_metal_fusion_check_norm },
- { GGML_METAL_FUSION_NORM_SCALE, ops_rms_norm_scale, {}, false, ggml_metal_fusion_check_norm },
- { GGML_METAL_FUSION_ADD_CHAIN, ops_add_2, {}, false, ggml_metal_fusion_check_add_chain },
- { GGML_METAL_FUSION_ADD_CHAIN, ops_add_3, {}, false, ggml_metal_fusion_check_add_chain },
- { GGML_METAL_FUSION_ADD_CHAIN, ops_add_4, {}, false, ggml_metal_fusion_check_add_chain },
- { GGML_METAL_FUSION_ADD_CHAIN, ops_add_5, {}, false, ggml_metal_fusion_check_add_chain },
- { GGML_METAL_FUSION_ADD_CHAIN, ops_add_6, {}, false, ggml_metal_fusion_check_add_chain },
- { GGML_METAL_FUSION_ADD_CHAIN, ops_add_7, {}, false, ggml_metal_fusion_check_add_chain },
- { GGML_METAL_FUSION_SNAKE, ops_snake, {}, false, ggml_metal_fusion_check_snake },
- { GGML_METAL_FUSION_GDN_CACHE, ops_gdn_cache, {}, true, ggml_metal_fusion_check_gdn_cache },
- { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe, {1}, true, ggml_metal_fusion_check_topk_moe },
- { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_scale, {1}, true, ggml_metal_fusion_check_topk_moe },
- { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_norm, {1}, true, ggml_metal_fusion_check_topk_moe },
- { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_norm_scale, {1}, true, ggml_metal_fusion_check_topk_moe },
- { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_2, {}, true, ggml_metal_fusion_check_moe_reduce },
- { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_3, {}, true, ggml_metal_fusion_check_moe_reduce },
- { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_4, {}, true, ggml_metal_fusion_check_moe_reduce },
- { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_5, {}, true, ggml_metal_fusion_check_moe_reduce },
- { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_6, {}, true, ggml_metal_fusion_check_moe_reduce },
- { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_7, {}, true, ggml_metal_fusion_check_moe_reduce },
- { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_8, {}, true, ggml_metal_fusion_check_moe_reduce },
- { GGML_METAL_FUSION_SSM_CONV_SILU, ops_ssm_conv_silu, {}, false, ggml_metal_fusion_check_ssm_conv_silu },
- { GGML_METAL_FUSION_MUL_MAT_ADD, ops_mul_mat_add, {}, false, ggml_metal_fusion_check_mul_mat_add },
+ { GGML_METAL_FUSION_NORM_MUL, ops_norm_mul, {}, false, ggml_metal_fusion_check_norm },
+ { GGML_METAL_FUSION_NORM_MUL_ADD, ops_norm_mul_add, {}, false, ggml_metal_fusion_check_norm },
+ { GGML_METAL_FUSION_NORM_SCALE, ops_norm_scale, {}, false, ggml_metal_fusion_check_norm },
+ { GGML_METAL_FUSION_NORM_MUL, ops_rms_norm_mul, {}, false, ggml_metal_fusion_check_norm },
+ { GGML_METAL_FUSION_NORM_MUL_ADD, ops_rms_norm_mul_add, {}, false, ggml_metal_fusion_check_norm },
+ { GGML_METAL_FUSION_NORM_SCALE, ops_rms_norm_scale, {}, false, ggml_metal_fusion_check_norm },
+
+ { GGML_METAL_FUSION_ADD_CHAIN, ops_add_2, {}, false, ggml_metal_fusion_check_add_chain },
+ { GGML_METAL_FUSION_ADD_CHAIN, ops_add_3, {}, false, ggml_metal_fusion_check_add_chain },
+ { GGML_METAL_FUSION_ADD_CHAIN, ops_add_4, {}, false, ggml_metal_fusion_check_add_chain },
+ { GGML_METAL_FUSION_ADD_CHAIN, ops_add_5, {}, false, ggml_metal_fusion_check_add_chain },
+ { GGML_METAL_FUSION_ADD_CHAIN, ops_add_6, {}, false, ggml_metal_fusion_check_add_chain },
+ { GGML_METAL_FUSION_ADD_CHAIN, ops_add_7, {}, false, ggml_metal_fusion_check_add_chain },
+
+ { GGML_METAL_FUSION_SNAKE, ops_snake, {}, false, ggml_metal_fusion_check_snake },
+
+ { GGML_METAL_FUSION_GDN_CACHE, ops_gdn_cache, {}, true, ggml_metal_fusion_check_gdn_cache },
+
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe, {1}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_scale, {1}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_norm, {1}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_norm_scale, {1}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_norm_bias, {1}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_norm_scale_bias, {1}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_sp, {3}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_sp_scale, {3}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_sp_norm, {3}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_sp_norm_scale, {3}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_sp_norm_bias, {3}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_sp_norm_scale_bias, {3}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_sig, {2}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_sig_scale, {2}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_sig_norm, {2}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_sig_norm_scale, {2}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_sig_norm_bias, {2}, true, ggml_metal_fusion_check_topk_moe },
+ { GGML_METAL_FUSION_TOPK_MOE, ops_topk_moe_sig_norm_scale_bias, {2}, true, ggml_metal_fusion_check_topk_moe },
+
+ { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_2, {}, true, ggml_metal_fusion_check_moe_reduce },
+ { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_3, {}, true, ggml_metal_fusion_check_moe_reduce },
+ { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_4, {}, true, ggml_metal_fusion_check_moe_reduce },
+ { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_5, {}, true, ggml_metal_fusion_check_moe_reduce },
+ { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_6, {}, true, ggml_metal_fusion_check_moe_reduce },
+ { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_7, {}, true, ggml_metal_fusion_check_moe_reduce },
+ { GGML_METAL_FUSION_MOE_REDUCE, ops_moe_reduce_8, {}, true, ggml_metal_fusion_check_moe_reduce },
+
+ { GGML_METAL_FUSION_SSM_CONV_SILU, ops_ssm_conv_silu, {}, false, ggml_metal_fusion_check_ssm_conv_silu },
+
+ { GGML_METAL_FUSION_MUL_MAT_ADD, ops_mul_mat_add, {}, false, ggml_metal_fusion_check_mul_mat_add },
};
// ---- alloc deps -----------------------------------------------------------
diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h
index 11c15feff..5fc49f9bb 100644
--- a/ggml/src/ggml-metal/ggml-metal-impl.h
+++ b/ggml/src/ggml-metal/ggml-metal-impl.h
@@ -1270,6 +1270,7 @@ typedef struct {
uint64_t nb01; // logits row stride
uint64_t nb1_ids; // ids row stride
float clamp;
+ float bias;
float scale;
} ggml_metal_kargs_topk_moe;
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index cace34bd4..a42f66080 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -821,6 +821,16 @@ int ggml_metal_op_acc(ggml_metal_op_t ctx, int idx) {
int ggml_metal_op_unary(ggml_metal_op_t ctx, int idx) {
ggml_tensor * op = ctx->node(idx);
+ // the sigmoid and sqrt(softplus) TOPK_MOE fusions anchor at the gating unary node
+ if (op->op == GGML_OP_UNARY && ctx->use_fusion() &&
+ (ggml_get_unary_op(op) == GGML_UNARY_OP_SOFTPLUS || ggml_get_unary_op(op) == GGML_UNARY_OP_SIGMOID)) {
+ int n = 1;
+ const ggml_metal_fusion * fusion = ctx->can_fuse(idx, GGML_METAL_FUSION_FULL, &n);
+ if (fusion && ggml_metal_fusion_get_id(fusion) == GGML_METAL_FUSION_TOPK_MOE) {
+ return ggml_metal_op_topk_moe(ctx, idx);
+ }
+ }
+
ggml_metal_library_t lib = ctx->lib;
ggml_metal_encoder_t enc = ctx->enc;
@@ -6023,29 +6033,48 @@ int ggml_metal_op_topk_moe(ggml_metal_op_t ctx, int idx) {
return 1;
}
- ggml_tensor * softmax = ctx->node(idx);
- ggml_tensor * logits = softmax->src[0];
- ggml_tensor * get_rows = ctx->node(idx + 2);
- ggml_tensor * ids = get_rows->src[1];
- ggml_tensor * weights = ctx->node(idx + n_fuse - 1);
+ ggml_tensor * first = ctx->node(idx);
+ ggml_tensor * logits = first->src[0];
+
+ // the unary chains score with sigmoid or sqrt(softplus) and rank biased scores
+ const bool with_unary = first->op == GGML_OP_UNARY;
+ const int gating = !with_unary ? 0 : ggml_get_unary_op(first) == GGML_UNARY_OP_SOFTPLUS ? 1 : 2;
+
+ ggml_tensor * get_rows = nullptr;
+ ggml_tensor * sum_rows = nullptr;
+ ggml_tensor * prob_bias = nullptr;
+ int sum_rows_pos = -1;
+ for (int j = 0; j < n_fuse; ++j) {
+ ggml_tensor * node = ctx->node(idx + j);
+ if (node->op == GGML_OP_GET_ROWS) { get_rows = node; }
+ if (node->op == GGML_OP_ADD) { prob_bias = node->src[1]; }
+ if (node->op == GGML_OP_SUM_ROWS) { sum_rows = node; sum_rows_pos = j; }
+ }
+
+ ggml_tensor * ids = get_rows->src[1];
+ ggml_tensor * weights = ctx->node(idx + n_fuse - 1);
+
+ const bool with_norm = sum_rows != nullptr;
+ const bool with_scale = weights->op == GGML_OP_SCALE;
const int64_t n_expert = logits->ne[0];
const int64_t n_tokens = logits->ne[1];
const int64_t n_expert_used = ids->ne[0];
- const bool with_norm = n_fuse >= 6;
- const bool with_scale = n_fuse == 4 || n_fuse == 7;
-
float clamp = -INFINITY;
+ float bias = 0.0f;
if (with_norm) {
- ggml_tensor * clamp_node = ctx->node(idx + 4);
- clamp = ggml_get_op_params_f32(clamp_node, 0);
+ ggml_tensor * denom_node = ctx->node(idx + sum_rows_pos + 1);
+ if (denom_node->op == GGML_OP_SCALE) {
+ bias = ggml_get_op_params_f32(denom_node, 1);
+ } else {
+ clamp = ggml_get_op_params_f32(denom_node, 0);
+ }
}
float scale = 1.0f;
if (with_scale) {
- ggml_tensor * scale_node = ctx->node(idx + n_fuse - 1);
- scale = ggml_get_op_params_f32(scale_node, 0);
+ scale = ggml_get_op_params_f32(weights, 0);
}
ggml_metal_kargs_topk_moe args = {
@@ -6053,23 +6082,27 @@ int ggml_metal_op_topk_moe(ggml_metal_op_t ctx, int idx) {
/*.nb01 =*/ logits->nb[1],
/*.nb1_ids =*/ ids->nb[1],
/*.clamp =*/ clamp,
+ /*.bias =*/ bias,
/*.scale =*/ scale,
};
- auto pipeline = ggml_metal_library_get_pipeline_topk_moe(lib, (int32_t) n_expert, (int32_t) n_expert_used, with_norm);
+ auto pipeline = ggml_metal_library_get_pipeline_topk_moe(lib, (int32_t) n_expert, (int32_t) n_expert_used,
+ with_norm, gating, prob_bias != nullptr);
ggml_metal_encoder_set_pipeline(enc, pipeline);
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(logits), 1);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(weights), 2);
ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(ids), 3);
+ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(prob_bias ? prob_bias : logits), 4);
ggml_metal_encoder_dispatch_threadgroups(enc, (uint32_t) n_tokens, 1, 1, 32, 1, 1);
ctx->count_fusions(fusion);
if (ggml_metal_fusion_info_debug(ctx->finfo) > 1) {
- GGML_LOG_DEBUG("%s: fuse: SOFT_MAX + ARGSORT + GET_ROWS\n", __func__);
+ GGML_LOG_DEBUG("%s: fuse: %s + ARGSORT + GET_ROWS\n", __func__,
+ gating == 1 ? "SOFTPLUS + SQRT" : gating == 2 ? "SIGMOID" : "SOFT_MAX");
}
return n_fuse;
diff --git a/ggml/src/ggml-metal/kernels/argsort.metal b/ggml/src/ggml-metal/kernels/argsort.metal
index 5231b8395..6328528a9 100644
--- a/ggml/src/ggml-metal/kernels/argsort.metal
+++ b/ggml/src/ggml-metal/kernels/argsort.metal
@@ -3,6 +3,8 @@
constant bool FC_topk_moe_with_norm [[function_constant(FC_TOPK_MOE + 0)]];
constant int FC_topk_moe_n_expert [[function_constant(FC_TOPK_MOE + 1)]];
constant int FC_topk_moe_top_k [[function_constant(FC_TOPK_MOE + 2)]];
+constant int FC_topk_moe_gating [[function_constant(FC_TOPK_MOE + 3)]];
+constant bool FC_topk_moe_pbias [[function_constant(FC_TOPK_MOE + 4)]];
constant int FC_moe_reduce_n_expert_used [[function_constant(FC_MOE_REDUCE + 0)]];
@@ -342,13 +344,14 @@ kernel void kernel_top_k_f32_i32(
}
}
-// fused SOFT_MAX + top-k + GET_ROWS (+ optional norm/scale) for MoE routing.
+// fused SOFT_MAX or sqrt(softplus) + top-k + GET_ROWS (+ optional norm/scale) for MoE routing.
// One SIMDgroup handles one token row; n_expert is limited to 1024 by the host.
kernel void kernel_topk_moe_f32(
constant ggml_metal_kargs_topk_moe & args,
device const char * src0,
device float * weights,
device int32_t * ids,
+ device const float * probs_bias,
uint3 tgpig[[threadgroup_position_in_grid]],
ushort tiisg[[thread_index_in_simdgroup]]) {
const int row = (int) tgpig.x;
@@ -377,50 +380,78 @@ kernel void kernel_topk_moe_f32(
wt[i / 32] = isnan(v) ? -FLT_MAX : v;
}
- // softmax over the expert logits
- float max_val = -INFINITY;
- FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
- max_val = max(max_val, wt[i]);
- }
- max_val = simd_max(max_val);
+ if (FC_topk_moe_gating == 1) {
+ FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
+ const float v = wt[i];
+ wt[i] = sqrt(v > 20.0f ? v : log(1.0f + exp(v)));
+ }
+ } else if (FC_topk_moe_gating == 2) {
+ FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
+ wt[i] = 1.0f / (1.0f + exp(-wt[i]));
+ }
+ } else {
+ // softmax over the expert logits
+ float max_val = -INFINITY;
+ FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
+ max_val = max(max_val, wt[i]);
+ }
+ max_val = simd_max(max_val);
- float sum_val = 0.0f;
- FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
- wt[i] = exp(wt[i] - max_val);
- sum_val += wt[i];
+ float sum_val = 0.0f;
+ FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
+ wt[i] = exp(wt[i] - max_val);
+ sum_val += wt[i];
+ }
+ sum_val = simd_sum(sum_val);
+
+ const float inv_sum = 1.0f / sum_val;
+ FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
+ wt[i] *= inv_sum;
+ }
}
- sum_val = simd_sum(sum_val);
- const float inv_sum = 1.0f / sum_val;
- FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
- wt[i] *= inv_sum;
+ // the ranking keys add the selection bias, the output weights stay unbiased
+ float sel[32];
+ if (FC_topk_moe_pbias) {
+ FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
+ const int expert = lane + i * 32;
+ sel[i] = expert < n_expert ? wt[i] + probs_bias[expert] : -INFINITY;
+ }
}
float wt_sum = 0.0f;
for (int k = 0; k < top_k; ++k) {
- float best_val = -INFINITY;
+ float best_val = -INFINITY;
+ float best_key = -INFINITY;
int best_expert = -1;
FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
- const int expert = lane + i * 32;
- if (expert < n_expert && (wt[i] > best_val || (wt[i] == best_val && expert < best_expert))) {
+ const int expert = lane + i * 32;
+ const float key = FC_topk_moe_pbias ? sel[i] : wt[i];
+ if (expert < n_expert && (key > best_key || (key == best_key && expert < best_expert))) {
best_val = wt[i];
+ best_key = key;
best_expert = expert;
}
}
FOR_UNROLL (int mask = 16; mask > 0; mask >>= 1) {
const float val = simd_shuffle_xor(best_val, mask);
+ const float key = simd_shuffle_xor(best_key, mask);
const int expert = simd_shuffle_xor(best_expert, mask);
- if (val > best_val || (val == best_val && expert < best_expert)) {
+ if (key > best_key || (key == best_key && expert < best_expert)) {
best_val = val;
+ best_key = key;
best_expert = expert;
}
}
if ((best_expert & 31) == lane) {
wt[best_expert / 32] = -INFINITY;
+ if (FC_topk_moe_pbias) {
+ sel[best_expert / 32] = -INFINITY;
+ }
}
if ((k & 31) == lane) {
@@ -437,7 +468,7 @@ kernel void kernel_topk_moe_f32(
if (FC_topk_moe_with_norm) {
wt_sum = simd_sum(wt_sum);
- wt_sum = max(wt_sum, args.clamp);
+ wt_sum = max(wt_sum, args.clamp) + args.bias;
const float inv = 1.0f / wt_sum;
FOR_UNROLL (int i = 0; i < n_per_lane; ++i) {
output_weights[i] *= inv;
diff --git a/tests/fusion/MTL.csv b/tests/fusion/MTL.csv
index eda5119b5..2ce291c7d 100644
--- a/tests/fusion/MTL.csv
+++ b/tests/fusion/MTL.csv
@@ -3,6 +3,7 @@
afmoe ,1 ,any ,MUL+ADD , 2
afmoe ,1 ,any ,RMS_NORM+MUL , 10
afmoe ,1 ,any ,RMS_NORM+MUL+ADD , 3
+afmoe ,1 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
arcee ,0 ,any ,RMS_NORM+MUL , 5
arctic ,0 ,any ,MUL+ADD , 4
arctic ,0 ,any ,RMS_NORM+MUL , 7
@@ -15,11 +16,13 @@ bailingmoe ,1 ,any ,SOFT_MAX+ARGSORT+GET_ROWS , 2
bailingmoe2 ,1 ,any ,ADD+ADD , 1
bailingmoe2 ,1 ,any ,MUL+ADD , 2
bailingmoe2 ,1 ,any ,RMS_NORM+MUL , 9
+bailingmoe2 ,1 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
bailingmoe3 ,1 ,any ,ADD+ADD , 1
bailingmoe3 ,1 ,any ,GATED_DELTA_NET+CPY , 1
bailingmoe3 ,1 ,any ,MUL+ADD , 2
bailingmoe3 ,1 ,any ,RMS_NORM+MUL , 8
bailingmoe3 ,1 ,any ,RMS_NORM+SCALE , 2
+bailingmoe3 ,1 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
beam ,0 ,any ,MUL+ADD , 6
beam ,0 ,any ,RMS_NORM+MUL , 25
bloom ,0 ,any ,NORM+MUL+ADD , 6
@@ -50,21 +53,27 @@ deepseek ,0 ,any ,SOFT_MAX+ARGSORT+GET_ROWS , 1
deepseek2 ,0 ,any ,ADD+ADD , 1
deepseek2 ,0 ,any ,MUL+ADD , 2
deepseek2 ,0 ,any ,RMS_NORM+MUL , 9
+deepseek2 ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
deepseek32 ,0 ,any ,ADD+ADD , 1
deepseek32 ,0 ,any ,MUL+ADD , 2
deepseek32 ,0 ,any ,NORM+MUL+ADD , 2
deepseek32 ,0 ,any ,RMS_NORM+MUL , 9
+deepseek32 ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
deepseek4 ,0 ,any ,MUL+ADD , 8
deepseek4 ,0 ,any ,RMS_NORM+MUL , 20
+deepseek4 ,0 ,any ,UNARY+SQRT+ADD+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV, 4
deepseek41 ,0 ,any ,MUL+ADD , 12
deepseek41 ,0 ,any ,RMS_NORM+MUL , 29
+deepseek41 ,0 ,any ,UNARY+SQRT+ADD+ARGSORT+GET_ROWS+SUM_ROWS+SCALE+DIV, 6
dots1 ,0 ,any ,ADD+ADD , 1
dots1 ,0 ,any ,MUL+ADD , 2
dots1 ,0 ,any ,RMS_NORM+MUL , 9
+dots1 ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
dots3note ,0 ,any ,ADD+ADD , 1
dots3note ,0 ,any ,MUL+ADD , 2
dots3note ,0 ,any ,NORM+MUL+ADD , 1
dots3note ,0 ,any ,RMS_NORM+MUL , 11
+dots3note ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
dream ,0 ,any ,RMS_NORM+MUL , 5
ernie4_5-moe ,1 ,any ,ADD+ADD , 1
ernie4_5-moe ,1 ,any ,MUL+ADD , 2
@@ -74,6 +83,7 @@ exaone ,0 ,any ,RMS_NORM+MUL , 5
exaone-moe ,1 ,any ,ADD+ADD , 1
exaone-moe ,1 ,any ,MUL+ADD , 2
exaone-moe ,1 ,any ,RMS_NORM+MUL , 9
+exaone-moe ,1 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
exaone4 ,0 ,any ,RMS_NORM+MUL , 5
exaone4 ,0 ,any ,RMS_NORM+MUL+ADD , 4
falcon ,0 ,any ,ADD+ADD , 2
@@ -83,6 +93,7 @@ falcon-h1 ,0 ,any ,RMS_NORM+MUL , 9
flint ,0 ,any ,ADD+ADD , 1
flint ,0 ,any ,MUL+ADD , 2
flint ,0 ,any ,RMS_NORM+MUL , 9
+flint ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV, 1
gemma ,0 ,any ,RMS_NORM+MUL , 5
gemma2 ,0 ,any ,RMS_NORM+MUL , 5
gemma2 ,0 ,any ,RMS_NORM+MUL+ADD , 4
@@ -94,17 +105,20 @@ glm-dsa ,0 ,any ,ADD+ADD , 1
glm-dsa ,0 ,any ,MUL+ADD , 2
glm-dsa ,0 ,any ,NORM+MUL+ADD , 2
glm-dsa ,0 ,any ,RMS_NORM+MUL , 9
+glm-dsa ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
glm4 ,0 ,any ,RMS_NORM+MUL , 5
glm4 ,0 ,any ,RMS_NORM+MUL+ADD , 4
glm4moe ,1 ,any ,ADD+ADD , 1
glm4moe ,1 ,any ,MUL+ADD , 2
glm4moe ,1 ,any ,RMS_NORM+MUL , 9
+glm4moe ,1 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
glm5-next ,0 ,any ,ADD+ADD+ADD , 1
glm5-next ,0 ,any ,GATED_DELTA_NET+CPY , 1
glm5-next ,0 ,any ,MUL+ADD , 2
glm5-next ,0 ,any ,NORM+MUL+ADD , 1
glm5-next ,0 ,any ,RMS_NORM+MUL , 8
glm5-next ,0 ,any ,RMS_NORM+SCALE , 2
+glm5-next ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
gpt-oss ,0 ,any ,MUL+ADD , 4
gpt-oss ,0 ,any ,RMS_NORM+MUL , 5
gpt2 ,0 ,any ,NORM+MUL+ADD , 5
@@ -135,9 +149,11 @@ hunyuan_vl ,0 ,any ,RMS_NORM+MUL , 9
hy_v3 ,0 ,any ,ADD+ADD , 2
hy_v3 ,0 ,any ,MUL+ADD , 4
hy_v3 ,0 ,any ,RMS_NORM+MUL , 9
+hy_v3 ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 2
hy_v4 ,0 ,any ,MUL+ADD , 2
hy_v4 ,0 ,any ,NORM+MUL+ADD , 1
hy_v4 ,0 ,any ,RMS_NORM+MUL , 9
+hy_v4 ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV, 1
internlm2 ,0 ,any ,RMS_NORM+MUL , 5
jais ,0 ,any ,NORM+MUL+ADD , 5
jais2 ,0 ,any ,NORM+MUL+ADD , 5
@@ -146,24 +162,29 @@ k2-horizon ,0 ,any ,RMS_NORM+MUL , 4
k2-horizon ,0 ,any ,ADD+ADD , 1
k2-horizon ,0 ,any ,MUL+ADD , 2
k2-horizon ,0 ,any ,RMS_NORM+MUL , 4
+k2-horizon ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
kimi-k3 ,0 ,any ,GATED_DELTA_NET+CPY , 1
kimi-k3 ,0 ,any ,MUL+ADD , 2
kimi-k3 ,0 ,any ,RMS_NORM+MUL , 17
kimi-k3 ,0 ,any ,RMS_NORM+SCALE , 2
+kimi-k3 ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
kimi-linear ,0 ,any ,ADD+ADD , 1
kimi-linear ,0 ,any ,GATED_DELTA_NET+CPY , 1
kimi-linear ,0 ,any ,MUL+ADD , 2
kimi-linear ,0 ,any ,RMS_NORM+MUL , 7
kimi-linear ,0 ,any ,RMS_NORM+SCALE , 2
+kimi-linear ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV, 1
kolibri1 ,0 ,any ,MUL+ADD , 4
kolibri1 ,0 ,any ,RMS_NORM+MUL , 10
kolibri1 ,0 ,any ,RMS_NORM+MUL+ADD , 3
laguna ,0 ,any ,ADD+ADD , 1
laguna ,0 ,any ,MUL+ADD , 2
laguna ,0 ,any ,RMS_NORM+MUL , 9
+laguna ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
lfm2 ,0 ,any ,RMS_NORM+MUL , 7
lfm2moe ,1 ,any ,MUL+ADD , 2
lfm2moe ,1 ,any ,RMS_NORM+MUL , 7
+lfm2moe ,1 ,any ,UNARY+ADD+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV, 1
lfm3 ,0 ,any ,RMS_NORM+MUL , 17
llada ,0 ,any ,RMS_NORM+MUL , 5
llada-moe ,1 ,any ,MUL+ADD , 4
@@ -192,6 +213,7 @@ metis16 ,0 ,any ,MUL+ADD , 8
metis16 ,0 ,any ,RMS_NORM+MUL , 11
mimo2 ,0 ,any ,MUL+ADD , 4
mimo2 ,0 ,any ,RMS_NORM+MUL , 5
+mimo2 ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV, 2
minicpm ,0 ,any ,RMS_NORM+MUL , 5
minicpm ,0 ,any ,MUL+ADD , 4
minicpm ,0 ,any ,RMS_NORM+MUL , 5
@@ -202,9 +224,11 @@ minimax-01 ,0 ,any ,RMS_NORM+MUL , 6
minimax-01 ,0 ,any ,SOFT_MAX+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV, 2
minimax-m2 ,0 ,any ,MUL+ADD , 4
minimax-m2 ,0 ,any ,RMS_NORM+MUL , 9
+minimax-m2 ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV, 2
minimax-m3 ,0 ,any ,ADD+ADD , 1
minimax-m3 ,0 ,any ,MUL+ADD , 2
minimax-m3 ,0 ,any ,RMS_NORM+MUL , 11
+minimax-m3 ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
mistral-large4 ,0 ,any ,ADD+ADD , 3
mistral-large4 ,0 ,any ,MUL+ADD , 6
mistral-large4 ,0 ,any ,NORM+MUL+ADD , 6
@@ -216,11 +240,14 @@ mistral3 ,0 ,any ,SOFT_MAX+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV,
mistral4 ,0 ,any ,ADD+ADD , 1
mistral4 ,0 ,any ,MUL+ADD , 2
mistral4 ,0 ,any ,RMS_NORM+MUL , 9
+mistral4 ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS , 1
mpt ,0 ,any ,NORM+MUL+ADD , 5
muse-glimmer ,0 ,any ,RMS_NORM+MUL , 10
muse-glimmer ,0 ,any ,RMS_NORM+MUL+ADD , 3
-muse-spark ,0 ,any ,MUL+ADD , 4
+muse-spark ,0 ,prefill ,MUL+ADD , 2
+muse-spark ,0 ,decode ,MUL+ADD , 4
muse-spark ,0 ,any ,RMS_NORM+MUL , 14
+muse-spark ,0 ,any ,UNARY+ADD+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV, 2
nanbeige ,0 ,any ,RMS_NORM+MUL , 5
nemotron ,0 ,any ,NORM+MUL+ADD , 5
nemotron_h ,0 ,any ,RMS_NORM+MUL , 5
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index c63b5a571..e35806f31 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -7551,6 +7551,7 @@ struct test_topk_moe : public test_case {
const int n_expert_groups;
ggml_tensor * logits {};
ggml_tensor * exp_probs_b {};
+ const bool norm_bias;
ggml_tensor * weights {};
ggml_tensor * selected_experts {};
@@ -7560,20 +7561,23 @@ struct test_topk_moe : public test_case {
bool bias_probs = false,
MoeGatingFunc gating_func = GATING_FUNC_SOFTMAX,
float scale_w = 0.0f,
- int n_expert_groups = 1) :
+ int n_expert_groups = 1,
+ bool norm_bias = false) :
ne(ne),
n_expert_used(n_expert_used),
with_norm(with_norm),
bias_probs(bias_probs),
gating_func(gating_func),
scale_w(scale_w),
- n_expert_groups(n_expert_groups) {
+ n_expert_groups(n_expert_groups),
+ norm_bias(norm_bias) {
GGML_ASSERT(n_expert_used <= ne[0]);
GGML_ASSERT(ne[0] % n_expert_groups == 0);
GGML_ASSERT(n_expert_used <= ne[0] / n_expert_groups);
+ GGML_ASSERT(!norm_bias || with_norm);
}
- std::string vars() override { return VARS_TO_STR7(ne, n_expert_used, with_norm, bias_probs, gating_func, scale_w, n_expert_groups); }
+ std::string vars() override { return VARS_TO_STR8(ne, n_expert_used, with_norm, norm_bias, bias_probs, gating_func, scale_w, n_expert_groups); }
std::string op_desc(ggml_tensor * t) override {
GGML_UNUSED(t);
@@ -7731,7 +7735,11 @@ struct test_topk_moe : public test_case {
ggml_tensor * weights_sum = ggml_sum_rows(ctx, weights); // [1, n_tokens]
ggml_set_name(weights_sum, "weights_sum");
- weights_sum = ggml_clamp(ctx, weights_sum, 6.103515625e-5, INFINITY);
+ if (norm_bias) {
+ weights_sum = ggml_scale_bias(ctx, weights_sum, 1.0f, 1e-20f);
+ } else {
+ weights_sum = ggml_clamp(ctx, weights_sum, 6.103515625e-5, INFINITY);
+ }
weights = ggml_div(ctx, weights, weights_sum); // [n_expert_used, n_tokens]
weights = ggml_reshape_3d(ctx, weights, 1, n_expert_used, n_tokens);
}
@@ -12871,6 +12879,22 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
}
+ // the sum is adjusted with an additive bias instead of a lower clamp (DeepSeek-V4.1)
+ for (bool bias_probs : {false, true}) {
+ for (float scale_w : {0.0f, 2.0f}) {
+ test_cases.emplace_back(new test_topk_moe({8, 22, 1, 1}, 4, true, bias_probs, GATING_FUNC_SOFTMAX, scale_w, 1, true));
+ test_cases.emplace_back(new test_topk_moe({384, 1, 1, 1}, 6, true, bias_probs, GATING_FUNC_SOFTMAX, scale_w, 1, true));
+ test_cases.emplace_back(new test_topk_moe({8, 22, 1, 1}, 4, true, bias_probs, GATING_FUNC_SQRT_SOFTPLUS, scale_w, 1, true));
+ test_cases.emplace_back(new test_topk_moe({384, 1, 1, 1}, 6, true, bias_probs, GATING_FUNC_SQRT_SOFTPLUS, scale_w, 1, true));
+ }
+ }
+
+ // sigmoid scoring always ranks a selection bias (GLM, DeepSeek-V2/V3.2)
+ for (float scale_w : {0.0f, 2.0f}) {
+ test_cases.emplace_back(new test_topk_moe({8, 22, 1, 1}, 4, true, true, GATING_FUNC_SIGMOID, scale_w, 1, true));
+ test_cases.emplace_back(new test_topk_moe({384, 1, 1, 1}, 6, true, true, GATING_FUNC_SIGMOID, scale_w, 1, true));
+ }
+
// Cover the supported boundaries, common k = 8 shapes, interleaved views and adds, and k = 16 fallback.
test_cases.emplace_back(new test_moe_reduce(63, 2, 17));
test_cases.emplace_back(new test_moe_reduce(2048, 8, 128));