WIP

Author: roberteg16

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -4779,11 +4779,6 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph
     // store {fork_idx, join_idx}
     std::vector<std::pair<int, int>> concurrent_node_ranges;
 
-    // per-event lists of concurrent branch nodes to place in the dedicated scratch buffer (below),
-    // so the branches are mutually disjoint and disjoint from tensors read across the region -
-    // this replaces the fragile node interleaving that ggml-alloc/execution order can desync
-    std::vector<std::vector<const ggml_tensor *>> concurrent_groups;
-
     for (const auto & [root_node, count] : fan_out) {
         if (count >= min_fan_out && count <= max_fan_out) {
             const int root_node_idx = node_indices[root_node];
@@ -4875,6 +4870,10 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph
                 int fork_node_idx = node_indices[root_node];
                 int join_node_idx = node_indices[join_node];
 
+                int       current_branch_idx = 0;
+                int       current_node_idx   = fork_node_idx + 1;
+                const int n_branches         = nodes_per_branch.size();
+
                 int total_branch_nodes = 0;
                 for (std::vector<const ggml_tensor *> branch_nodes : nodes_per_branch) {
                     total_branch_nodes += branch_nodes.size();
@@ -4903,20 +4902,48 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph
                 GGML_LOG_DEBUG("Adding stream at node %s %p\n", root_node->name, root_node);
                 concurrent_node_ranges.emplace_back(fork_node_idx, join_node_idx);
 
-                // place all branch nodes in the dedicated scratch buffer (below) instead of
-                // interleaving them: ggml-alloc then keeps the branches mutually disjoint and
-                // disjoint from the fork output that every branch reads concurrently
-                std::vector<const ggml_tensor *> group;
-                for (const auto & branch_nodes : nodes_per_branch) {
-                    for (const ggml_tensor * n : branch_nodes) {
-                        group.push_back(n);
+                // interleave tensors to extend lifetimes so that ggml graph doesn't recycle them
+                // example transformation:
+                // [attn-norm, QMul, QNorm, QRope, KMul, KNorm, KRope, VMul, attn] ->
+                // [attn-norm, QMul, KMul, VMul, QNorm, VNorm, QRope, KRope, attn]
+                while (current_node_idx < join_node_idx) {
+                    std::vector<const ggml_tensor *> & branch_nodes = nodes_per_branch[current_branch_idx];
+
+                    bool has_node = false;
+                    for (std::vector<const ggml_tensor *> branch_node : nodes_per_branch) {
+                        has_node |= branch_node.size() > 0;
+                    }
+
+                    GGML_ASSERT(has_node);
+
+                    if (branch_nodes.empty()) {
+                        current_branch_idx = (current_branch_idx + 1) % n_branches;
+                        continue;
+                    }
+
+                    cgraph->nodes[current_node_idx] = const_cast<ggml_tensor *>(branch_nodes.front());
+                    current_node_idx++;
+                    branch_nodes.erase(branch_nodes.begin());
+
+                    // append all empty nodes
+                    while (!branch_nodes.empty() && is_noop(branch_nodes.front())) {
+                        cgraph->nodes[current_node_idx] = const_cast<ggml_tensor *>(branch_nodes.front());
+                        current_node_idx++;
+                        branch_nodes.erase(branch_nodes.begin());
                     }
+
+                    current_branch_idx = (current_branch_idx + 1) % n_branches;
                 }
-                concurrent_groups.push_back(std::move(group));
             }
         }
     }
 
+    // per-event lists of concurrent branch nodes to place in the dedicated scratch buffer (below),
+    // so the branches are mutually disjoint and disjoint from tensors read across the region.
+    // The QKV path below uses the legacy node interleaving instead; only the MoE shared-expert path
+    // populates this, so the two disjointness mechanisms coexist.
+    std::vector<std::vector<const ggml_tensor *>> concurrent_groups;
+
     ggml_cuda_detect_shared_expert_concurrency(cgraph, cuda_ctx, node_indices, concurrent_node_ranges, concurrent_groups);
 
     // Place every concurrent branch (attention QKV and MoE shared-expert) in a dedicated buffer so