Skip to content

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

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

Open
jeffbolznv wants to merge 1 commit into
base: master
Choose a base branch
from
+85 −12
Open

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

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
89 changes: 77 additions & 12 deletions ggml/src/ggml-vulkan/ggml-vulkan.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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_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,
GGML_OP_VIEW, GGML_OP_GET_ROWS };

Expand Down Expand Up @@ -673,6 +680,22 @@ static constexpr std::initializer_list<std::array<int, 3>> topk_moe_sigmoid_norm
{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
static constexpr std::initializer_list<std::array<int, 3>> topk_moe_early_softmax_edges {
{ 1, 0, 0 }, // reshape->src[0] == softmax
Expand Down Expand Up @@ -701,6 +724,7 @@ enum topk_moe_mode {
TOPK_MOE_EARLY_SOFTMAX_NORM,
TOPK_MOE_LATE_SOFTMAX,
TOPK_MOE_SIGMOID_NORM_BIAS,
TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS,
TOPK_MOE_COUNT,
};

Expand Down Expand Up @@ -13105,12 +13129,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) {
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 * 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 * ids = (mode == TOPK_MOE_SIGMOID_NORM_BIAS) ? cgraph->nodes[node_idx + 4] :
(mode == TOPK_MOE_LATE_SOFTMAX) ? cgraph->nodes[node_idx + 1] :
cgraph->nodes[node_idx + 3];
ggml_tensor * ids = mode == TOPK_MOE_SIGMOID_NORM_BIAS ? cgraph->nodes[node_idx + 4] :
mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS ? cgraph->nodes[node_idx + 5] :
mode == TOPK_MOE_LATE_SOFTMAX ? cgraph->nodes[node_idx + 1] :
cgraph->nodes[node_idx + 3];

GGML_ASSERT(logits->type == GGML_TYPE_F32);
GGML_ASSERT(bias->type == GGML_TYPE_F32);
Expand Down Expand Up @@ -13150,16 +13178,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_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_SIGMOID 1
#define GATING_FUNC_SOFTMAX_WEIGHT 2

pc.gating_func = mode == TOPK_MOE_SIGMOID_NORM_BIAS ? GATING_FUNC_SIGMOID :
mode == TOPK_MOE_LATE_SOFTMAX ? GATING_FUNC_SOFTMAX_WEIGHT :
GATING_FUNC_SOFTMAX;
pc.has_bias = mode == TOPK_MOE_SIGMOID_NORM_BIAS;
pc.with_norm = mode == TOPK_MOE_EARLY_SOFTMAX_NORM || mode == TOPK_MOE_SIGMOID_NORM_BIAS;
#define GATING_FUNC_SQRT_SOFTPLUS 3

pc.gating_func = mode == TOPK_MOE_SIGMOID_NORM_BIAS ? GATING_FUNC_SIGMOID :
mode == TOPK_MOE_SQRT_SOFTPLUS_NORM_BIAS ? GATING_FUNC_SQRT_SOFTPLUS :
mode == TOPK_MOE_LATE_SOFTMAX ? GATING_FUNC_SOFTMAX_WEIGHT :
GATING_FUNC_SOFTMAX;
pc.has_bias = has_bias;
pc.with_norm = mode == TOPK_MOE_EARLY_SOFTMAX_NORM || has_bias;
if (ctx->fused_topk_moe_scale) {
GGML_ASSERT(weights->op == GGML_OP_SCALE);
pc.output_scale = ggml_get_op_params_f32(weights, 0);
Expand Down Expand Up @@ -16241,6 +16277,20 @@ static bool ggml_vk_can_fuse_topk_moe(ggml_backend_vk_context * ctx, const struc
return false;
}
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:
softmax = cgraph->nodes[node_idx + 0];
weights = cgraph->nodes[node_idx + 4];
Expand All @@ -16264,7 +16314,9 @@ static bool ggml_vk_can_fuse_topk_moe(ggml_backend_vk_context * ctx, const struc
probs = probs->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;
}

Expand Down Expand Up @@ -16632,7 +16684,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 src is allowed to overlap the memory for the destination.
// 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_scale = false;
Expand Down Expand Up @@ -16743,6 +16795,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;
fusion_string = "TOPK_MOE_SIGMOID_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_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 }) &&
ggml_check_edges(cgraph, i, topk_moe_early_softmax_edges) &&
ggml_vk_can_fuse_topk_moe(ctx, cgraph, i, TOPK_MOE_EARLY_SOFTMAX)) {
Expand Down Expand Up @@ -17009,6 +17070,9 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
if (keep_pattern(topk_moe_sigmoid_norm_bias)) {
continue;
}
if (keep_pattern(topk_moe_sqrt_softplus_norm_bias)) {
continue;
}
if (keep_pattern(topk_moe_early_softmax)) {
continue;
}
Expand Down Expand Up @@ -17039,6 +17103,7 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
// Don't pull forward nodes from fusion patterns
if (match_pattern(topk_moe_early_softmax_norm, 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_late_softmax, j) ||
match_pattern(snake_pattern, j)) {
Expand Down
8 changes: 8 additions & 0 deletions ggml/src/ggml-vulkan/vulkan-shaders/topk_moe.comp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#define GATING_FUNC_SOFTMAX 0
#define GATING_FUNC_SIGMOID 1
#define GATING_FUNC_SOFTMAX_WEIGHT 2
#define GATING_FUNC_SQRT_SOFTPLUS 3

layout (push_constant) uniform parameter
{
Expand Down Expand Up @@ -120,6 +121,13 @@ void main() {
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;
}
} 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];
Expand Down
Loading