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

vulkan: extend topk_moe fusion to support sqrt(softplus) (llama/26124)

This commit is contained in:
Jeff Bolz 2026-08-01 14:18:07 -05:00 committed by Georgi Gerganov
parent cb625868e8
commit 83105b7c3c
2 changed files with 84 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

87
ggml/src/ggml-vulkan/ggml-vulkan.cpp
View File

@ -610,6 +610,13 @@ static constexpr std::initializer_list<ggml_op> topk_moe_sigmoid_norm_bias{ GGML
GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP,
GGML_OP_DIV, GGML_OP_RESHAPE }; GGML_OP_DIV, GGML_OP_RESHAPE };
static constexpr std::initializer_list<ggml_op> topk_moe_sqrt_softplus_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_CLAMP,
GGML_OP_DIV, GGML_OP_RESHAPE };
static constexpr std::initializer_list<ggml_op> topk_moe_early_softmax { GGML_OP_SOFT_MAX, GGML_OP_RESHAPE, GGML_OP_ARGSORT, static constexpr std::initializer_list<ggml_op> topk_moe_early_softmax { GGML_OP_SOFT_MAX, GGML_OP_RESHAPE, GGML_OP_ARGSORT,
GGML_OP_VIEW, GGML_OP_GET_ROWS }; GGML_OP_VIEW, GGML_OP_GET_ROWS };
@ -673,6 +680,22 @@ static constexpr std::initializer_list<std::array<int, 3>> topk_moe_sigmoid_norm
{10, 0, 9 }, // reshape->src[0] == div {10, 0, 9 }, // reshape->src[0] == div
}; };
static constexpr std::initializer_list<std::array<int, 3>> topk_moe_sqrt_softplus_norm_bias_edges {
{ 1, 0, 0 }, // sqrt->src[0] == softplus
{ 2, 0, 1 }, // reshape->src[0] == sqrt
{ 3, 0, 1 }, // add->src[0] == sqrt
{ 4, 0, 3 }, // argsort->src[0] == add
{ 5, 0, 4 }, // view->src[0] == argsort
{ 6, 0, 2 }, // get_rows->src[0] == reshape
{ 6, 1, 5 }, // get_rows->src[1] == view
{ 7, 0, 6 }, // reshape->src[0] == get_rows
{ 8, 0, 7 }, // sum_rows->src[0] == reshape
{ 9, 0, 8 }, // clamp->src[0] == sum_rows
{10, 0, 7 }, // div->src[0] == reshape
{10, 1, 9 }, // div->src[1] == clamp
{11, 0,10 }, // reshape->src[0] == div
};
// same as early_softmax_norm but ending after the get_rows // same as early_softmax_norm but ending after the get_rows
static constexpr std::initializer_list<std::array<int, 3>> topk_moe_early_softmax_edges { static constexpr std::initializer_list<std::array<int, 3>> topk_moe_early_softmax_edges {
{ 1, 0, 0 }, // reshape->src[0] == softmax { 1, 0, 0 }, // reshape->src[0] == softmax
@ -701,6 +724,7 @@ enum topk_moe_mode {
TOPK_MOE_EARLY_SOFTMAX_NORM, TOPK_MOE_EARLY_SOFTMAX_NORM,
TOPK_MOE_LATE_SOFTMAX, TOPK_MOE_LATE_SOFTMAX,
TOPK_MOE_SIGMOID_NORM_BIAS, TOPK_MOE_SIGMOID_NORM_BIAS,
TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS,
TOPK_MOE_COUNT, TOPK_MOE_COUNT,
}; };
@ -13203,12 +13227,16 @@ static void ggml_vk_soft_max_back(ggml_backend_vk_context * ctx, vk_context& sub
static void ggml_vk_topk_moe(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx) { static void ggml_vk_topk_moe(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_cgraph * cgraph, int node_idx) {
topk_moe_mode mode = ctx->fused_topk_moe_mode; topk_moe_mode mode = ctx->fused_topk_moe_mode;
const bool has_bias = mode == TOPK_MOE_SIGMOID_NORM_BIAS || mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS;
ggml_tensor * logits = cgraph->nodes[node_idx + 0]->src[0]; ggml_tensor * logits = cgraph->nodes[node_idx + 0]->src[0];
ggml_tensor * bias = (mode == TOPK_MOE_SIGMOID_NORM_BIAS) ? cgraph->nodes[node_idx + 2]->src[1] : logits; ggml_tensor * bias = mode == TOPK_MOE_SIGMOID_NORM_BIAS ? cgraph->nodes[node_idx + 2]->src[1] :
mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS ? cgraph->nodes[node_idx + 3]->src[1] :
logits;
ggml_tensor * weights = cgraph->nodes[node_idx + ctx->num_additional_fused_ops]; ggml_tensor * weights = cgraph->nodes[node_idx + ctx->num_additional_fused_ops];
ggml_tensor * ids = (mode == TOPK_MOE_SIGMOID_NORM_BIAS) ? cgraph->nodes[node_idx + 4] : ggml_tensor * ids = mode == TOPK_MOE_SIGMOID_NORM_BIAS ? cgraph->nodes[node_idx + 4] :
(mode == TOPK_MOE_LATE_SOFTMAX) ? cgraph->nodes[node_idx + 1] : mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS ? cgraph->nodes[node_idx + 5] :
cgraph->nodes[node_idx + 3]; mode == TOPK_MOE_LATE_SOFTMAX ? cgraph->nodes[node_idx + 1] :
cgraph->nodes[node_idx + 3];
GGML_ASSERT(logits->type == GGML_TYPE_F32); GGML_ASSERT(logits->type == GGML_TYPE_F32);
GGML_ASSERT(bias->type == GGML_TYPE_F32); GGML_ASSERT(bias->type == GGML_TYPE_F32);
@ -13248,16 +13276,24 @@ static void ggml_vk_topk_moe(ggml_backend_vk_context * ctx, vk_context& subctx,
pc.clamp_min = ggml_get_op_params_f32(clamp, 0); pc.clamp_min = ggml_get_op_params_f32(clamp, 0);
pc.clamp_max = ggml_get_op_params_f32(clamp, 1); pc.clamp_max = ggml_get_op_params_f32(clamp, 1);
} }
if (mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS) {
ggml_tensor * clamp = cgraph->nodes[node_idx + 9];
GGML_ASSERT(clamp->op == GGML_OP_CLAMP);
pc.clamp_min = ggml_get_op_params_f32(clamp, 0);
pc.clamp_max = ggml_get_op_params_f32(clamp, 1);
}
#define GATING_FUNC_SOFTMAX 0 #define GATING_FUNC_SOFTMAX 0
#define GATING_FUNC_SIGMOID 1 #define GATING_FUNC_SIGMOID 1
#define GATING_FUNC_SOFTMAX_WEIGHT 2 #define GATING_FUNC_SOFTMAX_WEIGHT 2
#define GATING_FUNC_SQRT_SOFTPLUS 3
pc.gating_func = mode == TOPK_MOE_SIGMOID_NORM_BIAS ? GATING_FUNC_SIGMOID : pc.gating_func = mode == TOPK_MOE_SIGMOID_NORM_BIAS ? GATING_FUNC_SIGMOID :
mode == TOPK_MOE_LATE_SOFTMAX ? GATING_FUNC_SOFTMAX_WEIGHT : mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS ? GATING_FUNC_SQRT_SOFTPLUS :
GATING_FUNC_SOFTMAX; mode == TOPK_MOE_LATE_SOFTMAX ? GATING_FUNC_SOFTMAX_WEIGHT :
pc.has_bias = mode == TOPK_MOE_SIGMOID_NORM_BIAS; GATING_FUNC_SOFTMAX;
pc.with_norm = mode == TOPK_MOE_EARLY_SOFTMAX_NORM || mode == TOPK_MOE_SIGMOID_NORM_BIAS; pc.has_bias = has_bias;
pc.with_norm = mode == TOPK_MOE_EARLY_SOFTMAX_NORM || has_bias;
if (ctx->fused_topk_moe_scale) { if (ctx->fused_topk_moe_scale) {
GGML_ASSERT(weights->op == GGML_OP_SCALE); GGML_ASSERT(weights->op == GGML_OP_SCALE);
pc.output_scale = ggml_get_op_params_f32(weights, 0); pc.output_scale = ggml_get_op_params_f32(weights, 0);
@ -16366,6 +16402,20 @@ static bool ggml_vk_can_fuse_topk_moe(ggml_backend_vk_context * ctx, const struc
return false; return false;
} }
break; break;
case TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS:
softmax = cgraph->nodes[node_idx + 0]; // really softplus
weights = cgraph->nodes[node_idx + 11];
get_rows = cgraph->nodes[node_idx + 6];
argsort = cgraph->nodes[node_idx + 4];
if (ggml_get_unary_op(softmax) != GGML_UNARY_OP_SOFTPLUS) {
return false;
}
// bias is expected to be 1D
if (ggml_nrows(cgraph->nodes[node_idx + 3]->src[1]) != 1 ||
!ggml_is_contiguous(cgraph->nodes[node_idx + 3]->src[1])) {
return false;
}
break;
case TOPK_MOE_EARLY_SOFTMAX: case TOPK_MOE_EARLY_SOFTMAX:
softmax = cgraph->nodes[node_idx + 0]; softmax = cgraph->nodes[node_idx + 0];
weights = cgraph->nodes[node_idx + 4]; weights = cgraph->nodes[node_idx + 4];
@ -16389,7 +16439,9 @@ static bool ggml_vk_can_fuse_topk_moe(ggml_backend_vk_context * ctx, const struc
probs = probs->src[0]; probs = probs->src[0];
ggml_tensor * selection_probs = argsort->src[0]; ggml_tensor * selection_probs = argsort->src[0];
if (probs != selection_probs && mode != TOPK_MOE_SIGMOID_NORM_BIAS) { if (probs != selection_probs &&
mode != TOPK_MOE_SIGMOID_NORM_BIAS &&
mode != TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS) {
return false; return false;
} }
@ -16757,7 +16809,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
// the fused result in an elementwise-way. This affects whether the memory for // the fused result in an elementwise-way. This affects whether the memory for
// the src is allowed to overlap the memory for the destination. // the src is allowed to overlap the memory for the destination.
// The array is sized to handle the largest fusion (asserted later). // The array is sized to handle the largest fusion (asserted later).
bool op_srcs_fused_elementwise[12]; bool op_srcs_fused_elementwise[13];
ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_mode = TOPK_MOE_COUNT;
ctx->fused_topk_moe_scale = false; ctx->fused_topk_moe_scale = false;
@ -16868,6 +16920,15 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
ctx->fused_topk_moe_mode = TOPK_MOE_SIGMOID_NORM_BIAS; ctx->fused_topk_moe_mode = TOPK_MOE_SIGMOID_NORM_BIAS;
fusion_string = "TOPK_MOE_SIGMOID_NORM_BIAS"; fusion_string = "TOPK_MOE_SIGMOID_NORM_BIAS";
std::fill_n(op_srcs_fused_elementwise, ctx->num_additional_fused_ops + 1, false); std::fill_n(op_srcs_fused_elementwise, ctx->num_additional_fused_ops + 1, false);
} else if (ggml_can_fuse_subgraph(cgraph, i, topk_moe_sqrt_softplus_norm_bias, { i + 5, i + 11 }) &&
ggml_check_edges(cgraph, i, topk_moe_sqrt_softplus_norm_bias_edges) &&
ggml_vk_can_fuse_topk_moe(ctx, cgraph, i, TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS)) {
ctx->num_additional_fused_ops = topk_moe_sqrt_softplus_norm_bias.size() - 1;
// view of argsort writes to memory
ctx->fused_ops_write_mask |= 1 << 5;
ctx->fused_topk_moe_mode = TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS;
fusion_string = "TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS";
std::fill_n(op_srcs_fused_elementwise, ctx->num_additional_fused_ops + 1, false);
} else if (ggml_can_fuse_subgraph(cgraph, i, topk_moe_early_softmax, { i + 3, i + 4 }) && } else if (ggml_can_fuse_subgraph(cgraph, i, topk_moe_early_softmax, { i + 3, i + 4 }) &&
ggml_check_edges(cgraph, i, topk_moe_early_softmax_edges) && ggml_check_edges(cgraph, i, topk_moe_early_softmax_edges) &&
ggml_vk_can_fuse_topk_moe(ctx, cgraph, i, TOPK_MOE_EARLY_SOFTMAX)) { ggml_vk_can_fuse_topk_moe(ctx, cgraph, i, TOPK_MOE_EARLY_SOFTMAX)) {
@ -17134,6 +17195,9 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
if (keep_pattern(topk_moe_sigmoid_norm_bias)) { if (keep_pattern(topk_moe_sigmoid_norm_bias)) {
continue; continue;
} }
if (keep_pattern(topk_moe_sqrt_softplus_norm_bias)) {
continue;
}
if (keep_pattern(topk_moe_early_softmax)) { if (keep_pattern(topk_moe_early_softmax)) {
continue; continue;
} }
@ -17164,6 +17228,7 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
// Don't pull forward nodes from fusion patterns // Don't pull forward nodes from fusion patterns
if (match_pattern(topk_moe_early_softmax_norm, j) || if (match_pattern(topk_moe_early_softmax_norm, j) ||
match_pattern(topk_moe_sigmoid_norm_bias, j) || match_pattern(topk_moe_sigmoid_norm_bias, j) ||
match_pattern(topk_moe_sqrt_softplus_norm_bias, j) ||
match_pattern(topk_moe_early_softmax, j) || match_pattern(topk_moe_early_softmax, j) ||
match_pattern(topk_moe_late_softmax, j) || match_pattern(topk_moe_late_softmax, j) ||
match_pattern(snake_pattern, j)) { match_pattern(snake_pattern, j)) {

8
ggml/src/ggml-vulkan/vulkan-shaders/topk_moe.comp
View File

@ -10,6 +10,7 @@
#define GATING_FUNC_SOFTMAX 0 #define GATING_FUNC_SOFTMAX 0
#define GATING_FUNC_SIGMOID 1 #define GATING_FUNC_SIGMOID 1
#define GATING_FUNC_SOFTMAX_WEIGHT 2 #define GATING_FUNC_SOFTMAX_WEIGHT 2
#define GATING_FUNC_SQRT_SOFTPLUS 3
layout (push_constant) uniform parameter layout (push_constant) uniform parameter
{ {
@ -120,6 +121,13 @@ void main() {
const uint expert = i + lane; const uint expert = i + lane;
probs[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? 1.f / (1.f + exp(-probs[i / WARP_SIZE])) : -INFINITY; probs[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? 1.f / (1.f + exp(-probs[i / WARP_SIZE])) : -INFINITY;
} }
} else if (gating_func == GATING_FUNC_SQRT_SOFTPLUS) {
[[unroll]]
for (uint i = 0; i < n_experts; i += WARP_SIZE) {
const uint expert = i + lane;
const float val = probs[i / WARP_SIZE];
probs[i / WARP_SIZE] = (n_experts % WARP_SIZE == 0 || expert < n_experts) ? sqrt(val > 20.0f ? val : log(1.0f + exp(val))) : -INFINITY;
}
} }
float selection_probs[experts_per_thread]; float selection_probs[experts_per_thread];