Commit 50182a53f for llama.cpp
commit 50182a53fa2c26bd2a7fc31d855231effdc2f4ad
Author: Jeff Bolz <jbolz@nvidia.com>
Date: Thu Sep 10 12:21:29 2026 -0500
vulkan: use add_alloc_dep to enable topk_moe fusion for prefill (#28422)
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index 34400a1b7..8e1cf3ff3 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -18328,39 +18328,31 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
bool need_disable = false;
- // topk_moe often overwrites the source, but for a given row all the src values are
- // loaded before anything is stored. If there's only one row, this is safe, so treat
- // this as a special case.
- bool is_topk_moe_single_row = ctx->fused_topk_moe_mode != TOPK_MOE_COUNT &&
- ggml_nrows(cgraph->nodes[i]->src[0]) == 1;
-
- if (!is_topk_moe_single_row) {
- for (int j = 0; j < 2; ++j) {
- ggml_tensor *dst = output_nodes[j];
- if (!dst) {
- continue;
- }
- // Loop over all srcs of all nodes in the fusion. If the src overlaps
- // the destination and the src is not an intermediate node that's being
- // elided, then disable fusion.
- for (int k = 0; k <= ctx->num_additional_fused_ops; ++k) {
- for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) {
- ggml_tensor *src = cgraph->nodes[i + k]->src[s];
- if (!src || src->op == GGML_OP_NONE) {
- continue;
- }
- if (ggml_vk_tensors_overlap(src, dst, op_srcs_fused_elementwise[k])) {
- bool found = false;
- for (int n = 0; n < k; ++n) {
- if (cgraph->nodes[i + n] == src) {
- found = true;
- break;
- }
- }
- if (!found) {
- need_disable = true;
+ for (int j = 0; j < 2; ++j) {
+ ggml_tensor *dst = output_nodes[j];
+ if (!dst) {
+ continue;
+ }
+ // Loop over all srcs of all nodes in the fusion. If the src overlaps
+ // the destination and the src is not an intermediate node that's being
+ // elided, then disable fusion.
+ for (int k = 0; k <= ctx->num_additional_fused_ops; ++k) {
+ for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) {
+ ggml_tensor *src = cgraph->nodes[i + k]->src[s];
+ if (!src || src->op == GGML_OP_NONE) {
+ continue;
+ }
+ if (ggml_vk_tensors_overlap(src, dst, op_srcs_fused_elementwise[k])) {
+ bool found = false;
+ 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_qsa = false;
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.
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)");
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;
};
- 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;
}
- if (keep_pattern(topk_moe_sigmoid_norm_bias)) {
+ if (keep_topk_moe_pattern(topk_moe_sigmoid_norm_bias)) {
continue;
}
- if (keep_pattern(topk_moe_sqrt_softplus_norm_bias)) {
+ if (keep_topk_moe_pattern(topk_moe_sqrt_softplus_norm_bias)) {
continue;
}
- if (keep_pattern(topk_moe_early_softmax)) {
+ if (keep_topk_moe_pattern(topk_moe_early_softmax)) {
continue;
}
- if (keep_pattern(topk_moe_late_softmax)) {
+ if (keep_topk_moe_pattern(topk_moe_late_softmax)) {
continue;
}
if (keep_pattern(snake_pattern)) {