upstream-archive/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-07-23 11:10:55 -05:00
Code Issues 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

42
ggml/src/ggml-cuda/ggml-cuda.cu
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
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,
GGML_OP_GET_ROWS });
out_nodes[0] = i + 4;
ids = cgraph->nodes[i + 4];
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(),
{ 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_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_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]];
if (args.norm) {
ops.insert(ops.end(),

16
ggml/src/ggml-cuda/topk-moe.cu
View File

@@ -8,6 +8,7 @@
// Kernel config struct - passed by value to CUDA kernel
struct topk_moe_config {
bool use_sigmoid;
bool use_sqrt_softplus;
bool with_norm;
bool delayed_softmax;
};
@@ -67,6 +68,16 @@ __device__ void sigmoid_warp_inplace(float (&vals)[experts_per_thread], const in
}
}
template <int experts_per_thread, bool use_limit>
__device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], const int limit, const int lane) {
#pragma unroll
for (int i = 0; i < experts_per_thread; i++) {
const int idx = lane + i * WARP_SIZE;
const bool active = !use_limit || (idx < limit);
vals[i] = active ? sqrtf(vals[i] > 20.0f ? vals[i] : logf(1.0f + expf(vals[i]))) : -INFINITY;
}
}
/*
This kernel does the following:
1. optionally softmax over the logits per token [n_experts, n_tokens]
@@ -115,6 +126,8 @@ __launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float *
if (!config.delayed_softmax) {
if (config.use_sigmoid) {
sigmoid_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
} else if (config.use_sqrt_softplus) {
sqrt_softplus_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
} else {
softmax_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
}
@@ -365,6 +378,7 @@ void ggml_cuda_op_topk_moe(ggml_backend_cuda_context & ctx,
topk_moe_config config;
config.use_sigmoid = args.sigmoid;
config.use_sqrt_softplus = args.sqrt_softplus;
config.with_norm = with_norm;
config.delayed_softmax = args.delayed_softmax;
@@ -415,7 +429,7 @@ bool ggml_cuda_should_use_topk_moe(const ggml_tensor * gating_op,
} else if (gating_op->op == GGML_OP_UNARY) {
ggml_unary_op op = ggml_get_unary_op(gating_op);
if (op != GGML_UNARY_OP_SIGMOID) {
if (op != GGML_UNARY_OP_SIGMOID && op != GGML_UNARY_OP_SOFTPLUS) {
return false;
}
}

1
ggml/src/ggml-cuda/topk-moe.cuh
View File

@@ -5,6 +5,7 @@
struct ggml_cuda_topk_moe_args {
bool sigmoid{};
bool sqrt_softplus{};
bool softmax{};
bool delayed_softmax{};
bool prob_bias{};

7
tests/test-backend-ops.cpp
View File

@@ -5960,6 +5960,7 @@ enum MoeGatingFunc {
GATING_FUNC_SOFTMAX,
GATING_FUNC_SIGMOID,
GATING_FUNC_SOFTMAX_WEIGHT,
GATING_FUNC_SQRT_SOFTPLUS,
};
struct test_topk_moe : public test_case {
@@ -6003,7 +6004,8 @@ struct test_topk_moe : public test_case {
ggml_tensor * logits = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data());
ggml_tensor * probs =
(gating_func == GATING_FUNC_SOFTMAX) ? ggml_soft_max(ctx, logits) :
(gating_func == GATING_FUNC_SIGMOID) ? ggml_sigmoid(ctx, logits) : logits;
(gating_func == GATING_FUNC_SIGMOID) ? ggml_sigmoid(ctx, logits) :
(gating_func == GATING_FUNC_SQRT_SOFTPLUS) ? ggml_sqrt(ctx, ggml_softplus(ctx, logits)) : logits;
ggml_set_name(probs, "probs");
ggml_tensor * selection_probs = probs;
@@ -9584,7 +9586,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
}
for (auto gate : {GATING_FUNC_SOFTMAX, GATING_FUNC_SIGMOID, GATING_FUNC_SOFTMAX_WEIGHT}) {
for (auto gate : {GATING_FUNC_SOFTMAX, GATING_FUNC_SIGMOID, GATING_FUNC_SOFTMAX_WEIGHT, GATING_FUNC_SQRT_SOFTPLUS}) {
for (bool with_norm : {false, true}) {
for (bool bias_probs : {false, true}) {
for (float scale_w : {0.0f, 2.0f}) {
@@ -9596,6 +9598,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_topk_moe({128, 1, 1, 1}, 128, with_norm, bias_probs, gate, scale_w));
test_cases.emplace_back(new test_topk_moe({129, 1, 1, 1}, 128, with_norm, bias_probs, gate, scale_w));
test_cases.emplace_back(new test_topk_moe({160, 4, 1, 1}, 160, with_norm, bias_probs, gate, scale_w));
test_cases.emplace_back(new test_topk_moe({256, 22, 1, 1}, 6, with_norm, bias_probs, gate, scale_w)); // Used by DeepSeek-V4
test_cases.emplace_back(new test_topk_moe({288, 22, 1, 1}, 8, with_norm, bias_probs, gate, scale_w)); // Used by StepFun 3.7
}
}