Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

cuda: add sqrt_softplus in topk-moe for dsv4 (#25896)

Browse Source
This commit is contained in:
Aman Gupta
2026-07-22 00:30:01 +08:00
committed by GitHub
parent fb0e6b6219
commit 846e991ec3
4 changed files with 51 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/ggml-cuda.cu
+27 -13
View File
@@ -2703,6 +2703,7 @@ static int ggml_cuda_try_gdn_cache_fusion(
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;
args.softmax = false;
args.delayed_softmax = false;
args.prob_bias = false;
@@ -2716,10 +2717,17 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod
}
if (nodes[node_idx]->op == GGML_OP_UNARY) {
if (ggml_get_unary_op(nodes[node_idx]) != GGML_UNARY_OP_SIGMOID) {
const ggml_unary_op unary_op = ggml_get_unary_op(nodes[node_idx]);
if (unary_op == GGML_UNARY_OP_SIGMOID) {
args.sigmoid = true;
} else if (unary_op == GGML_UNARY_OP_SOFTPLUS && node_idx + 1 < n_nodes &&
nodes[node_idx + 1]->op == GGML_OP_SQRT && nodes[node_idx + 1]->src[0] == nodes[node_idx]) {
// sqrt(softplus(x)) scoring (DeepSeek-V4)
args.sqrt_softplus = true;
node_idx++;
} else {
return false;
}
args.sigmoid = true;
}
if (nodes[node_idx]->op == GGML_OP_ARGSORT) {
@@ -2728,7 +2736,7 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod
node_idx++;
if (args.sigmoid || args.softmax) {
if (args.sigmoid || args.sqrt_softplus || args.softmax) {
// SOFTMAX -> RESHAPE
if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_RESHAPE ||
nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
@@ -3172,21 +3180,27 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
const ggml_tensor * scale = nullptr;
if (!args.delayed_softmax) {
ggml_op gating_op = args.sigmoid ? GGML_OP_UNARY : GGML_OP_SOFT_MAX;
int out_nodes[2]; // nodes which can't be elided
int out_nodes[2]; // nodes which can't be elided
if (args.sigmoid) {
ops.insert(ops.end(), { GGML_OP_UNARY });
} else if (args.sqrt_softplus) {
ops.insert(ops.end(), { GGML_OP_UNARY, GGML_OP_SQRT });
} else {
ops.insert(ops.end(), { GGML_OP_SOFT_MAX });
}
const int i_probs = i + (int) ops.size() - 1; // last node of the gating activation
if (args.prob_bias) {
bias = cgraph->nodes[i + 2]->src[1];
ops.insert(ops.end(), { gating_op, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW,
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 + 4;
ids = cgraph->nodes[i + 4];
out_nodes[0] = i_probs + 4;
} else {
ops.insert(ops.end(),
{ gating_op, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS });
out_nodes[0] = i + 3;
ids = cgraph->nodes[i + 3];
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]];
if (args.norm) {
ops.insert(ops.end(),