Public/koboldcpp
1
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-09-20 01:31:42 +02:00
Code Issues Packages Projects Releases Wiki Activity

vulkan: use add_alloc_dep to enable topk_moe fusion for prefill (#28422)

Browse Source
This commit is contained in:
Jeff Bolz
2026-09-10 12:21:29 -05:00
committed by GitHub
parent 6788edb4f3
commit 50182a53fa
1 changed files with 60 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-vulkan/ggml-vulkan.cpp
+60 -37
View File
@@ -18328,38 +18328,30 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
bool need_disable = false; bool need_disable = false;
// topk_moe often overwrites the source, but for a given row all the src values are for (int j = 0; j < 2; ++j) {
// loaded before anything is stored. If there's only one row, this is safe, so treat ggml_tensor *dst = output_nodes[j];
// this as a special case. if (!dst) {
bool is_topk_moe_single_row = ctx->fused_topk_moe_mode != TOPK_MOE_COUNT && continue;
ggml_nrows(cgraph->nodes[i]->src[0]) == 1; }
// Loop over all srcs of all nodes in the fusion. If the src overlaps
if (!is_topk_moe_single_row) { // the destination and the src is not an intermediate node that's being
for (int j = 0; j < 2; ++j) { // elided, then disable fusion.
ggml_tensor *dst = output_nodes[j]; for (int k = 0; k <= ctx->num_additional_fused_ops; ++k) {
if (!dst) { for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) {
continue; ggml_tensor *src = cgraph->nodes[i + k]->src[s];
} if (!src || src->op == GGML_OP_NONE) {
// Loop over all srcs of all nodes in the fusion. If the src overlaps continue;
// the destination and the src is not an intermediate node that's being }
// elided, then disable fusion. if (ggml_vk_tensors_overlap(src, dst, op_srcs_fused_elementwise[k])) {
for (int k = 0; k <= ctx->num_additional_fused_ops; ++k) { bool found = false;
for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) { for (int n = 0; n < k; ++n) {
ggml_tensor *src = cgraph->nodes[i + k]->src[s]; if (cgraph->nodes[i + n] == src) {
if (!src || src->op == GGML_OP_NONE) { found = true;
continue; break;
}
} }
if (ggml_vk_tensors_overlap(src, dst, op_srcs_fused_elementwise[k])) { if (!found) {
bool found = false; need_disable = true;
for (int n = 0; n < k; ++n) {
if (cgraph->nodes[i + n] == src) {
found = true;
break;
}
}
if (!found) {
need_disable = true;
}
} }
} }
} }
@@ -18372,6 +18364,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
ctx->fused_topk_moe_scale = false; ctx->fused_topk_moe_scale = false;
ctx->fused_topk_qsa = false; ctx->fused_topk_qsa = false;
ctx->fused_rms_norm_mode = RMS_NORM_COUNT; ctx->fused_rms_norm_mode = RMS_NORM_COUNT;
fusion_string = nullptr;
} }
} }
@@ -18474,7 +18467,6 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
// Sort the graph for improved parallelism. // Sort the graph for improved parallelism.
static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * graph, struct ggml_backend_graph_optimize_params * params) static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * graph, struct ggml_backend_graph_optimize_params * params)
{ {
GGML_UNUSED(params);
VK_LOG_DEBUG("ggml_vk_graph_optimize(" << graph->n_nodes << " nodes)"); VK_LOG_DEBUG("ggml_vk_graph_optimize(" << graph->n_nodes << " nodes)");
ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context; ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context;
@@ -18560,19 +18552,50 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph *
return false; return false;
}; };
if (keep_pattern(topk_moe_early_softmax_norm)) { auto const &add_pattern_alloc_deps = [&](const std::initializer_list<ggml_op> &pattern, int last_node) {
// Keep external inputs alive through the fused output.
std::set<ggml_tensor *> seen;
for (size_t j = 0; j < pattern.size(); ++j) {
ggml_tensor * node = graph->nodes[first_unused + j];
for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) {
ggml_tensor * src = node->src[s];
if (src && seen.insert(src).second) {
params->add_alloc_dep(params->user_data, src, graph->nodes[last_node]);
}
}
seen.insert(node);
}
};
auto const &keep_topk_moe_pattern = [&](const std::initializer_list<ggml_op> &pattern) -> bool {
if (!match_pattern(pattern, first_unused)) {
return false;
}
int last_node = first_unused + (int) pattern.size() - 1;
// Some TOPK_MOE variants fuse a trailing scale.
if (last_node + 1 < graph->n_nodes && graph->nodes[last_node + 1]->op == GGML_OP_SCALE) {
last_node++;
}
add_pattern_alloc_deps(pattern, last_node);
return keep_pattern(pattern);
};
if (keep_topk_moe_pattern(topk_moe_early_softmax_norm)) {
continue; continue;
} }
if (keep_pattern(topk_moe_sigmoid_norm_bias)) { if (keep_topk_moe_pattern(topk_moe_sigmoid_norm_bias)) {
continue; continue;
} }
if (keep_pattern(topk_moe_sqrt_softplus_norm_bias)) { if (keep_topk_moe_pattern(topk_moe_sqrt_softplus_norm_bias)) {
continue; continue;
} }
if (keep_pattern(topk_moe_early_softmax)) { if (keep_topk_moe_pattern(topk_moe_early_softmax)) {
continue; continue;
} }
if (keep_pattern(topk_moe_late_softmax)) { if (keep_topk_moe_pattern(topk_moe_late_softmax)) {
continue; continue;
} }
if (keep_pattern(snake_pattern)) { if (keep_pattern(snake_pattern)) {