public_git
/
whisper.cpp
mirror of https://github.com/ggml-org/whisper.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

This commit is contained in:
Aman Gupta 2026-07-22 00:30:01 +08:00 committed by Georgi Gerganov
parent 812ad391d2
commit 7c9d38e5ae
No known key found for this signature in database
GPG Key ID: 449E073F9DC10735
3 changed files with 46 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

40
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) { 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.sigmoid = false;
args.sqrt_softplus = false;
args.softmax = false; args.softmax = false;
args.delayed_softmax = false; args.delayed_softmax = false;
args.prob_bias = 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 (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; return false;
} }
args.sigmoid = true;
} }
if (nodes[node_idx]->op == GGML_OP_ARGSORT) { 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++; node_idx++;
if (args.sigmoid || args.softmax) { if (args.sigmoid || args.sqrt_softplus || args.softmax) {
// SOFTMAX -> RESHAPE // SOFTMAX -> RESHAPE
if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_RESHAPE || if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_RESHAPE ||
nodes[node_idx]->src[0] != nodes[node_idx - 1]) { 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; const ggml_tensor * scale = nullptr;
if (!args.delayed_softmax) { 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) { if (args.prob_bias) {
bias = cgraph->nodes[i + 2]->src[1]; bias = cgraph->nodes[i_probs + 2]->src[1];
ops.insert(ops.end(), { gating_op, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW,
GGML_OP_GET_ROWS }); GGML_OP_GET_ROWS });
out_nodes[0] = i + 4; out_nodes[0] = i_probs + 4;
ids = cgraph->nodes[i + 4];
} else { } else {
ops.insert(ops.end(), ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS });
{ gating_op, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS }); out_nodes[0] = i_probs + 3;
out_nodes[0] = i + 3;
ids = cgraph->nodes[i + 3];
} }
ids = cgraph->nodes[out_nodes[0]];
if (args.norm) { if (args.norm) {
ops.insert(ops.end(), ops.insert(ops.end(),

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

@ -8,6 +8,7 @@
// Kernel config struct - passed by value to CUDA kernel // Kernel config struct - passed by value to CUDA kernel
struct topk_moe_config { struct topk_moe_config {
bool use_sigmoid; bool use_sigmoid;
bool use_sqrt_softplus;
bool with_norm; bool with_norm;
bool delayed_softmax; 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: This kernel does the following:
1. optionally softmax over the logits per token [n_experts, n_tokens] 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.delayed_softmax) {
if (config.use_sigmoid) { if (config.use_sigmoid) {
sigmoid_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x); 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 { } else {
softmax_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x); softmax_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
} }
@ -364,9 +377,10 @@ void ggml_cuda_op_topk_moe(ggml_backend_cuda_context & ctx,
} }
topk_moe_config config; topk_moe_config config;
config.use_sigmoid = args.sigmoid; config.use_sigmoid = args.sigmoid;
config.with_norm = with_norm; config.use_sqrt_softplus = args.sqrt_softplus;
config.delayed_softmax = args.delayed_softmax; config.with_norm = with_norm;
config.delayed_softmax = args.delayed_softmax;
if (bias) { 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, launch_topk_moe_cuda<true>(ctx, logits_d, weights_d, ids_d, bias_d, n_rows, n_experts, n_expert_used, clamp_val,
@ -415,7 +429,7 @@ bool ggml_cuda_should_use_topk_moe(const ggml_tensor * gating_op,
} else if (gating_op->op == GGML_OP_UNARY) { } else if (gating_op->op == GGML_OP_UNARY) {
ggml_unary_op op = ggml_get_unary_op(gating_op); 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; return false;
} }
} }

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

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