ggml-cuda: overlap the MoE shared expert on a separate stream

ggml/src/ggml-cuda/common.cuh

@@ -1244,10 +1244,6 @@ struct ggml_cuda_concurrent_event {
     int                                          n_streams = 0;
     std::unordered_map<const ggml_tensor *, int> stream_mapping;
 
-    // Original order of nodes in this concurrent region (before interleaving)
-    // Used to restore grouping for fusion within streams
-    std::vector<const ggml_tensor *> original_order;
-
     const ggml_tensor * join_node;
 
     ggml_cuda_concurrent_event() = default;
@@ -1270,7 +1266,6 @@ struct ggml_cuda_concurrent_event {
     , fork_event(other.fork_event)
     , n_streams(other.n_streams)
     , stream_mapping(std::move(other.stream_mapping))
-    , original_order(std::move(other.original_order))
     , join_node(other.join_node) {
         other.fork_event = nullptr;
     }
@@ -1394,7 +1389,9 @@ struct ggml_backend_cuda_context {
     cudaEvent_t copy_event = nullptr;
 
     cudaStream_t streams[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = { { nullptr } };
-    cublasHandle_t cublas_handles[GGML_CUDA_MAX_DEVICES] = {nullptr};
+    // one cuBLAS handle per (device, stream): the handle carries a workspace that must not be shared
+    // by concurrent streams, otherwise overlapped GEMMs corrupt each other's results
+    cublasHandle_t cublas_handles[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = { { nullptr } };
 
     int curr_stream_no = 0;
 
@@ -1457,6 +1454,11 @@ struct ggml_backend_cuda_context {
 
     ggml_cuda_stream_context concurrent_stream_context;
 
+    // dedicated buffer for the branches of overlapped concurrent regions (attention QKV, MoE shared
+    // expert), reused across layers so their scratch never aliases tensors read across the region
+    ggml_backend_buffer_t concurrent_scratch      = nullptr;
+    size_t                concurrent_scratch_size = 0;
+
     ~ggml_backend_cuda_context();
 
     cudaStream_t stream(int device, int stream) {
@@ -1472,12 +1474,13 @@ struct ggml_backend_cuda_context {
     ggml_cuda_stream_context & stream_context() { return concurrent_stream_context; }
 
     cublasHandle_t cublas_handle(int device) {
-        if (cublas_handles[device] == nullptr) {
+        cublasHandle_t & handle = cublas_handles[device][curr_stream_no];
+        if (handle == nullptr) {
             ggml_cuda_set_device(device);
-            CUBLAS_CHECK(cublasCreate(&cublas_handles[device]));
-            CUBLAS_CHECK(cublasSetMathMode(cublas_handles[device], CUBLAS_TF32_TENSOR_OP_MATH));
+            CUBLAS_CHECK(cublasCreate(&handle));
+            CUBLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TF32_TENSOR_OP_MATH));
         }
-        return cublas_handles[device];
+        return handle;
     }
 
     cublasHandle_t cublas_handle() {

ggml/src/ggml-cuda/concurrency.hpp

@@ -0,0 +1,331 @@
+#pragma once
+
+// Detection of the concurrent regions ("diamonds") that the CUDA backend overlaps
+// on separate streams: a fork node whose output fans into independent branches
+// that reconverge at a join node. Kept as a pure pass over the public ggml graph
+// API, with no CUDA state, so the heuristics can be exercised in isolation. The
+// backend turns each region into a ggml_cuda_concurrent_event and places its
+// branch tensors in a dedicated scratch buffer.
+
+#include "ggml.h"
+
+#include <algorithm>
+#include <cstring>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+struct ggml_cuda_concurrent_region {
+    int fork_idx = -1;
+    int join_idx = -1;
+
+    const ggml_tensor * fork_node = nullptr;
+    const ggml_tensor * join_node = nullptr;
+
+    int n_streams = 0;
+
+    // branch node -> 1-based aux stream number
+    std::unordered_map<const ggml_tensor *, int> stream_mapping;
+
+    // branch nodes to place in the dedicated scratch buffer, in graph order
+    std::vector<const ggml_tensor *> group;
+
+    // true for the MoE shared-expert diamond, false for the attention QKV fan-out
+    bool shared_expert = false;
+};
+
+inline std::vector<ggml_cuda_concurrent_region> ggml_cuda_detect_concurrent_regions(ggml_cgraph * cgraph) {
+    std::vector<ggml_cuda_concurrent_region> regions;
+
+    const int n_nodes = ggml_graph_n_nodes(cgraph);
+
+    const auto is_noop = [](const ggml_tensor * node) -> bool {
+        return ggml_is_empty(node) || node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE ||
+               node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE;
+    };
+
+    const auto depends_on = [](const ggml_tensor * dst, const ggml_tensor * src) -> bool {
+        for (int s = 0; s < GGML_MAX_SRC; ++s) {
+            if (dst->src[s] == src) {
+                return true;
+            }
+        }
+        // implicit dependency if they view the same tensor
+        const ggml_tensor * dst2 = dst->view_src ? dst->view_src : dst;
+        const ggml_tensor * src2 = src->view_src ? src->view_src : src;
+        return dst2 == src2;
+    };
+
+    // out-degree of each node (only counted from single-row consumers, see below)
+    std::unordered_map<const ggml_tensor *, int> fan_out;
+    // reverse mapping of node -> index in the graph
+    std::unordered_map<const ggml_tensor *, int> node_indices;
+
+    for (int node_idx = 0; node_idx < n_nodes; ++node_idx) {
+        const ggml_tensor * node = ggml_graph_node(cgraph, node_idx);
+        node_indices[node] = node_idx;
+
+        if (is_noop(node)) {
+            continue;
+        }
+        for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
+            const ggml_tensor * src = node->src[src_idx];
+            // only single-row (decode) consumers count toward fan-out
+            if (src && ggml_nrows(node) <= 1) {
+                fan_out[src] += 1;
+            }
+        }
+    }
+
+    // ranges [fork_idx, join_idx] already claimed by a region, to avoid nesting/overlap
+    std::vector<std::pair<int, int>> concurrent_node_ranges;
+    // fork nodes already used as a region key (dedup across attention and MoE passes)
+    std::unordered_set<const ggml_tensor *> used_forks;
+
+    // -----------------------------------------------------------------------------
+    // Attention QKV fan-out: a fork ("attn_norm") used by exactly N branches that
+    // reconverge at a join (e.g. flash-attn / KQ).
+    //   1. find fork nodes whose output is read by exactly N consumers
+    //   2. find the join, where >= 2 branch outputs are required
+    //   3. account for every branch node between fork and join
+    // -----------------------------------------------------------------------------
+    const int min_fan_out = 3;
+    const int max_fan_out = 3;
+
+    for (const auto & [root_node, count] : fan_out) {
+        if (count < min_fan_out || count > max_fan_out) {
+            continue;
+        }
+        const int root_node_idx = node_indices[root_node];
+
+        // only optimize for attn_norm
+        // TODO: make this more generic
+        if (!strstr(root_node->name, "attn_norm")) {
+            continue;
+        }
+
+        bool is_part_of_event = false;
+        for (const auto & [start, end] : concurrent_node_ranges) {
+            if (root_node_idx >= start && root_node_idx <= end) {
+                is_part_of_event = true;
+            }
+        }
+        if (is_part_of_event) {
+            continue;
+        }
+
+        std::vector<std::vector<const ggml_tensor *>> nodes_per_branch;
+        for (int i = root_node_idx + 1; i < n_nodes; ++i) {
+            const ggml_tensor * node = ggml_graph_node(cgraph, i);
+            if (!is_noop(node) && depends_on(node, root_node)) {
+                nodes_per_branch.push_back({ node });
+            }
+        }
+
+        GGML_ASSERT(nodes_per_branch.size() == (size_t) count);
+
+        // find the join point
+        const ggml_tensor * join_node = nullptr;
+
+        const auto & belongs_to_branch = [&](const ggml_tensor *                      node,
+                                             const std::vector<const ggml_tensor *> & branch) -> bool {
+            for (const ggml_tensor * n : branch) {
+                if (depends_on(node, n)) {
+                    return true;
+                }
+            }
+            return false;
+        };
+
+        for (int i = root_node_idx + 1; i < n_nodes; ++i) {
+            const ggml_tensor * curr_node = ggml_graph_node(cgraph, i);
+
+            int num_joins = 0;
+            for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) {
+                if (belongs_to_branch(curr_node, nodes_per_branch[branch_idx])) {
+                    num_joins++;
+                }
+            }
+
+            if (num_joins >= 2) {
+                join_node = curr_node;
+                break;
+            }
+
+            bool found_branch = false;
+            for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) {
+                std::vector<const ggml_tensor *> & branch_vec = nodes_per_branch[branch_idx];
+                if (belongs_to_branch(curr_node, branch_vec)) {
+                    if (std::find(branch_vec.begin(), branch_vec.end(), curr_node) == branch_vec.end()) {
+                        branch_vec.push_back(curr_node);
+                    }
+                    found_branch = true;
+                }
+            }
+
+            if (!found_branch && is_noop(curr_node)) {
+                // we can put it in any branch because it will be ignored
+                nodes_per_branch[0].push_back({ curr_node });
+            }
+        }
+
+        if (!join_node) {
+            continue;
+        }
+
+        const int fork_node_idx = root_node_idx;
+        const int join_node_idx = node_indices[join_node];
+
+        int total_branch_nodes = 0;
+        for (const auto & branch_nodes : nodes_per_branch) {
+            total_branch_nodes += (int) branch_nodes.size();
+        }
+
+        // there are other nodes in the middle which are unaccounted for
+        // (usually cpy nodes) - then ignore this fork
+        if (join_node_idx - fork_node_idx - 1 != total_branch_nodes) {
+            continue;
+        }
+
+        if (used_forks.count(root_node)) {
+            continue;
+        }
+
+        ggml_cuda_concurrent_region region;
+        region.fork_idx      = fork_node_idx;
+        region.join_idx      = join_node_idx;
+        region.fork_node     = root_node;
+        region.join_node     = join_node;
+        region.n_streams     = (int) nodes_per_branch.size();
+        region.shared_expert = false;
+
+        for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) {
+            for (const ggml_tensor * n : nodes_per_branch[branch_idx]) {
+                region.stream_mapping[n] = (int) branch_idx + 1;
+                region.group.push_back(n);
+            }
+        }
+
+        used_forks.insert(root_node);
+        concurrent_node_ranges.emplace_back(fork_node_idx, join_node_idx);
+        regions.push_back(std::move(region));
+    }
+
+    // -----------------------------------------------------------------------------
+    // MoE shared-expert overlap: run the shared expert on a separate stream,
+    // overlapped with the routed experts. fork = the FFN-input norm feeding both
+    // branches, join = ggml_add(ffn_moe_out, ffn_shexp*). Operands are matched by
+    // the names set via cb() in the model graph. Decode only (single row): prefill
+    // is compute-bound and gains nothing from the overlap.
+    // -----------------------------------------------------------------------------
+    const auto reach_backward = [](const ggml_tensor * start) {
+        std::unordered_set<const ggml_tensor *> seen;
+        std::vector<const ggml_tensor *> stack = { start };
+        while (!stack.empty()) {
+            const ggml_tensor * t = stack.back();
+            stack.pop_back();
+            if (!t || seen.count(t)) {
+                continue;
+            }
+            seen.insert(t);
+            for (int s = 0; s < GGML_MAX_SRC; ++s) {
+                if (t->src[s]) {
+                    stack.push_back(t->src[s]);
+                }
+            }
+        }
+        return seen;
+    };
+
+    for (int join_idx = 0; join_idx < n_nodes; ++join_idx) {
+        ggml_tensor * join_node = ggml_graph_node(cgraph, join_idx);
+        if (join_node->op != GGML_OP_ADD) {
+            continue;
+        }
+
+        // decode only: overlap helps only when the routed branch leaves the GPU
+        // underutilized (batch 1), not in prefill where matmuls already saturate it
+        if (ggml_nrows(join_node) > 1) {
+            continue;
+        }
+
+        ggml_tensor * routed_out = nullptr;
+        ggml_tensor * shexp_out  = nullptr;
+        for (int s = 0; s < 2; ++s) {
+            ggml_tensor * x = join_node->src[s];
+            ggml_tensor * y = join_node->src[1 - s];
+            if (x && y && strstr(x->name, "ffn_moe_out") && strstr(y->name, "ffn_shexp")) {
+                routed_out = x;
+                shexp_out  = y;
+            }
+        }
+        if (!routed_out || !shexp_out) {
+            continue;
+        }
+
+        const std::unordered_set<const ggml_tensor *> reach_routed = reach_backward(routed_out);
+        const std::unordered_set<const ggml_tensor *> reach_shexp  = reach_backward(shexp_out);
+
+        // fork = highest-index node reachable from both branches (the ffn_norm output)
+        int fork_idx = -1;
+        for (const ggml_tensor * t : reach_routed) {
+            if (!reach_shexp.count(t)) {
+                continue;
+            }
+            auto it = node_indices.find(t);
+            if (it != node_indices.end() && it->second < join_idx && it->second > fork_idx) {
+                fork_idx = it->second;
+            }
+        }
+        if (fork_idx < 0) {
+            continue;
+        }
+
+        bool overlaps = false;
+        for (const auto & [start, end] : concurrent_node_ranges) {
+            if (!(join_idx < start || fork_idx > end)) {
+                overlaps = true;
+            }
+        }
+        if (overlaps) {
+            continue;
+        }
+
+        // partition the region (fork_idx, join_idx): shared-expert nodes -> stream 1
+        std::vector<std::vector<const ggml_tensor *>> nodes_per_branch(2);
+        for (int i = fork_idx + 1; i < join_idx; ++i) {
+            const ggml_tensor * n = ggml_graph_node(cgraph, i);
+            const int branch = reach_shexp.count(n) ? 1 : 0;
+            nodes_per_branch[branch].push_back(n);
+        }
+        if (nodes_per_branch[0].empty() || nodes_per_branch[1].empty()) {
+            continue;
+        }
+
+        const ggml_tensor * fork_node = ggml_graph_node(cgraph, fork_idx);
+        if (used_forks.count(fork_node)) {
+            continue;
+        }
+
+        // the routed experts stay on the main stream and only the shared expert forks
+        // onto a single aux stream, joined at the add
+        ggml_cuda_concurrent_region region;
+        region.fork_idx      = fork_idx;
+        region.join_idx      = join_idx;
+        region.fork_node     = fork_node;
+        region.join_node     = join_node;
+        region.n_streams     = 1;
+        region.shared_expert = true;
+        for (const ggml_tensor * n : nodes_per_branch[1]) {
+            region.stream_mapping[n] = 1;
+            region.group.push_back(n);
+        }
+
+        used_forks.insert(fork_node);
+        concurrent_node_ranges.emplace_back(fork_idx, join_idx);
+        regions.push_back(std::move(region));
+    }
+
+    return regions;
+}