Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
am17an merged 1 commit into from
Jul 21, 2026
+51 −19
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 27 additions & 13 deletions ggml/src/ggml-cuda/ggml-cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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) {
Expand All @@ -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]) {
Expand Down Expand Up @@ -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(),
Expand Down
22 changes: 18 additions & 4 deletions ggml/src/ggml-cuda/topk-moe.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
};
Expand Down Expand Up @@ -67,6 +68,16 @@ __device__ void sigmoid_warp_inplace(float (&vals)[experts_per_thread], const in
}
}

template <int experts_per_thread, bool use_limit>

@ORippler ORippler Jul 21, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
template <int experts_per_thread, bool use_limit>
template <int experts_per_thread>

It seems like the function is instantiated with use_limit=false always, so we may remove this template parameter (other activations seem to be sharing this potentially -> may be worth follow-up PR)

__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;

@ORippler ORippler Jul 21, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
const int idx = lane + i * WARP_SIZE;
const int idx = lane + i * ggml_cuda_get_physical_warp_size();

This should resolve to 32 on CUDA and 64 on HIP afaik

@ORippler ORippler Jul 21, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Edit: Looking at the kernel-context, this is something that could/should be done in a separate follow-up PR for the whole op

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]
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -364,9 +377,10 @@ void ggml_cuda_op_topk_moe(ggml_backend_cuda_context & ctx,
}

topk_moe_config config;
config.use_sigmoid = args.sigmoid;
config.with_norm = with_norm;
config.delayed_softmax = args.delayed_softmax;
config.use_sigmoid = args.sigmoid;
config.use_sqrt_softplus = args.sqrt_softplus;
config.with_norm = with_norm;
config.delayed_softmax = args.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,
Expand Down Expand Up @@ -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;
}
}
Expand Down
1 change: 1 addition & 0 deletions ggml/src/ggml-cuda/topk-moe.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

struct ggml_cuda_topk_moe_args {
bool sigmoid{};
bool sqrt_softplus{};
bool softmax{};
Comment on lines 7 to 9

@ORippler ORippler Jul 21, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

May consider adding an enum in future, as sigmoid/sqrt_softplus and softmax seem to be mutually exclusive

bool delayed_softmax{};
bool prob_bias{};
Expand Down
7 changes: 5 additions & 2 deletions tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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}) {
Expand All @@ -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
}
}
Expand Down
Loading