From 0a3bbdbd64bb0da87fe6e8dad149991e8d3cb0fb Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Sun, 5 Jul 2026 00:04:22 -0700 Subject: [PATCH 1/6] ggml-cuda: honor the current stream in multi-stream concurrency Two correctness fixes for the GGML_CUDA_GRAPH_OPT fork/join concurrency so that work assigned to a non-default stream actually runs there: - ggml_cuda_op_mul_mat (the legacy tiled dispatcher) hardcoded ctx.stream(id, 0) for the src0 setup and ctx.stream(id, is) for the per-column compute loop. For a single GPU both resolve to stream 0 regardless of curr_stream_no, so any op on this path ran on the main stream while its aux-stream consumers only waited the fork event and read stale data. Use ctx.stream() for the non-split case; the multi-GPU split path keeps its per-device stream 0 for peer synchronization. This is a no-op when concurrency is off (curr_stream_no == 0). - The cuBLAS handle was shared across streams (one per device). The handle carries a workspace that concurrent GEMMs on different streams corrupt, so give each (device, stream) its own handle. Co-Authored-By: Claude Opus 4 (1M context) --- ggml/src/ggml-cuda/common.cuh | 13 ++++++++----- ggml/src/ggml-cuda/ggml-cuda.cu | 14 ++++++++++---- 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index e6e50e041195..8cd4f233cbd7 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1394,7 +1394,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; @@ -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() { diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index cca70592f807..c81e92343f65 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -614,8 +614,10 @@ ggml_backend_cuda_context::~ggml_backend_cuda_context() { CUDA_CHECK(cudaStreamDestroy(streams[i][j])); } } - if (cublas_handles[i] != nullptr) { - CUBLAS_CHECK(cublasDestroy(cublas_handles[i])); + for (int j = 0; j < GGML_CUDA_MAX_STREAMS; ++j) { + if (cublas_handles[i][j] != nullptr) { + CUBLAS_CHECK(cublasDestroy(cublas_handles[i][j])); + } } } } @@ -1925,7 +1927,9 @@ static void ggml_cuda_op_mul_mat( const bool dst_on_device = id == dst_ctx->device; ggml_cuda_set_device(id); - cudaStream_t stream = ctx.stream(id, 0); + // single-GPU work must run on the context's current stream so it participates correctly in + // multi-stream concurrency; multi-GPU split keeps its per-device stream 0 for peer sync + cudaStream_t stream = split ? ctx.stream(id, 0) : ctx.stream(); if (src0_is_contiguous) { dev[id].src0_dd = split ? (char *) src0_extra->data_device[id] : (char *) src0->data; @@ -1999,7 +2003,9 @@ static void ggml_cuda_op_mul_mat( const int64_t row_diff = dev[id].row_high - dev[id].row_low; ggml_cuda_set_device(id); - cudaStream_t stream = ctx.stream(id, is); + // single-GPU work must run on the context's current stream to participate in multi-stream + // concurrency; the per-column pipelining streams are only used by the multi-GPU split path + cudaStream_t stream = split ? ctx.stream(id, is) : ctx.stream(); // wait for main GPU data if necessary if (split && (id != ctx.device || is != 0)) { From e03ce76e906aff4254365b861e8b5ddea1c53800 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Sun, 5 Jul 2026 00:31:45 -0700 Subject: [PATCH 2/6] ggml-cuda: fix multi-stream concurrency by isolating branch scratch The attention QKV concurrency (GGML_CUDA_GRAPH_OPT) produced incorrect results because it relied on interleaving the branch nodes in the graph so that ggml-alloc would keep them non-overlapping. ggml-alloc sees the interleaved order while the executor restores the original order, and is_valid() only checks the branches against each other, not against tensors read across the region (e.g. the fork output every branch consumes concurrently). ggml-alloc could therefore recycle such a tensor's memory for a branch scratch, corrupting a concurrent read. Replace the interleave with a dedicated compute buffer for the concurrent branch nodes. Their data/buffer is pre-assigned into one CUDA buffer, reused across layers (which run sequentially) with distinct per-node offsets within a region, so the branches are structurally disjoint from each other and from the rest of the graph. is_valid() then accepts the event without any graph reordering. Co-Authored-By: Claude Opus 4 (1M context) --- ggml/src/ggml-cuda/common.cuh | 5 ++ ggml/src/ggml-cuda/ggml-cuda.cu | 89 ++++++++++++++++++++++----------- 2 files changed, 65 insertions(+), 29 deletions(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 8cd4f233cbd7..ffc01cc2ec79 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1459,6 +1459,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) { diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index c81e92343f65..4fde0ca29b5e 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -620,6 +620,9 @@ ggml_backend_cuda_context::~ggml_backend_cuda_context() { } } } + if (concurrent_scratch != nullptr) { + ggml_backend_buffer_free(concurrent_scratch); + } } @@ -4639,6 +4642,11 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph // store {fork_idx, join_idx} std::vector> 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> 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]; @@ -4730,10 +4738,6 @@ 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 branch_nodes : nodes_per_branch) { total_branch_nodes += branch_nodes.size(); @@ -4762,37 +4766,64 @@ 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); - // 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 & branch_nodes = nodes_per_branch[current_branch_idx]; - - bool has_node = false; - for (std::vector branch_node : nodes_per_branch) { - has_node |= branch_node.size() > 0; + // 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 group; + for (const auto & branch_nodes : nodes_per_branch) { + for (const ggml_tensor * n : branch_nodes) { + group.push_back(n); } + } + concurrent_groups.push_back(std::move(group)); + } + } + } - GGML_ASSERT(has_node); + // Place every concurrent branch in a dedicated buffer so its nodes never share an address with + // each other or with tensors read across the region (which ggml-alloc could otherwise recycle, + // corrupting concurrent reads). Layers run sequentially, so one buffer sized to the largest + // region is reused across all of them; within a region each node gets a distinct offset so the + // concurrent scratch stays disjoint. + if (!concurrent_groups.empty()) { + const size_t alignment = 128; - if (branch_nodes.empty()) { - current_branch_idx = (current_branch_idx + 1) % n_branches; - continue; - } + const auto group_footprint = [&](const std::vector & group) { + size_t off = 0; + for (const ggml_tensor * n : group) { + if (is_noop(n) || n->view_src != nullptr) { + continue; + } + off += GGML_PAD(ggml_nbytes(n), alignment); + } + return off; + }; + + size_t needed = 0; + for (const auto & group : concurrent_groups) { + needed = std::max(needed, group_footprint(group)); + } - cgraph->nodes[current_node_idx] = const_cast(branch_nodes.front()); - current_node_idx++; - branch_nodes.erase(branch_nodes.begin()); + if (needed > 0) { + if (cuda_ctx->concurrent_scratch == nullptr || cuda_ctx->concurrent_scratch_size < needed) { + if (cuda_ctx->concurrent_scratch != nullptr) { + ggml_backend_buffer_free(cuda_ctx->concurrent_scratch); + } + cuda_ctx->concurrent_scratch = ggml_backend_buft_alloc_buffer(ggml_backend_cuda_buffer_type(cuda_ctx->device), needed); + cuda_ctx->concurrent_scratch_size = needed; + } - // append all empty nodes - while (!branch_nodes.empty() && is_noop(branch_nodes.front())) { - cgraph->nodes[current_node_idx] = const_cast(branch_nodes.front()); - current_node_idx++; - branch_nodes.erase(branch_nodes.begin()); + char * const base = (char *) ggml_backend_buffer_get_base(cuda_ctx->concurrent_scratch); + for (const auto & group : concurrent_groups) { + size_t off = 0; + for (const ggml_tensor * cn : group) { + if (is_noop(cn) || cn->view_src != nullptr) { + continue; } - - current_branch_idx = (current_branch_idx + 1) % n_branches; + ggml_tensor * n = const_cast(cn); + n->data = base + off; + n->buffer = cuda_ctx->concurrent_scratch; + off += GGML_PAD(ggml_nbytes(n), alignment); } } } From df79627893137881a363c8b614cd5c9bda15db17 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Sun, 5 Jul 2026 00:32:13 -0700 Subject: [PATCH 3/6] ggml-cuda: overlap the MoE shared expert on a separate stream Add a shared-expert detector to ggml_backend_cuda_graph_optimize (after the attention pass, under GGML_CUDA_GRAPH_OPT). It matches the diamond join = ggml_add(ffn_moe_out, ffn_shexp*) by callback names and finds the FFN-input norm that forks into both branches. The routed experts stay on the main stream and only the shared expert forks onto a single aux stream, joined at the add. Keeping the large routed branch on the main stream (region nodes not in the stream map default to the main stream) avoids migrating it and needs just one fork/join. The shared-expert branch reuses the concurrent-branch scratch buffer so it is disjoint from the routed branch without any graph reordering. Gated to decode (single token). The overlap only helps when the routed branch leaves the GPU underutilized for the shared expert to run alongside it: that holds in decode (batch 1, latency/occupancy-bound) but not in prefill, where the routed matmuls already saturate the GPU and the overlap only adds contention. Output is byte-identical to sequential. Requires GGML_CUDA_GRAPH_OPT, CUDA graphs and a single GPU, so it is off by default. Co-Authored-By: Claude Opus 4 (1M context) --- ggml/src/ggml-cuda/ggml-cuda.cu | 139 ++++++++++++++++++++++++++++++-- 1 file changed, 131 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 4fde0ca29b5e..b9d654c0d8ce 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -85,6 +85,7 @@ #include #include #include +#include #include static_assert(sizeof(half) == sizeof(ggml_fp16_t), "wrong fp16 size"); @@ -4358,8 +4359,11 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud is_concurrent_event_active = false; concurrent_event = nullptr; } else { - GGML_ASSERT (concurrent_event->stream_mapping.find(node) != concurrent_event->stream_mapping.end()); - cuda_ctx->curr_stream_no = concurrent_event->stream_mapping[node]; + // region nodes not mapped to a concurrent stream run on the main stream: + // this keeps the routed branch on the main stream while only the shared + // expert forks off + auto it = concurrent_event->stream_mapping.find(node); + cuda_ctx->curr_stream_no = it != concurrent_event->stream_mapping.end() ? it->second : 0; GGML_LOG_DEBUG("Setting stream no to %d for node %s\n", cuda_ctx->curr_stream_no, node->name); } } else if (i - prev_i > 1) { @@ -4368,7 +4372,8 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud try_launch_concurrent_event(prev_node); if (is_concurrent_event_active) { - cuda_ctx->curr_stream_no = concurrent_event->stream_mapping[node]; + auto it = concurrent_event->stream_mapping.find(node); + cuda_ctx->curr_stream_no = it != concurrent_event->stream_mapping.end() ? it->second : 0; GGML_LOG_DEBUG("Setting stream no to %d for node %s\n", cuda_ctx->curr_stream_no, node->name); } } @@ -4780,11 +4785,129 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph } } - // Place every concurrent branch in a dedicated buffer so its nodes never share an address with - // each other or with tensors read across the region (which ggml-alloc could otherwise recycle, - // corrupting concurrent reads). Layers run sequentially, so one buffer sized to the largest - // region is reused across all of them; within a region each node gets a distinct offset so the - // concurrent scratch stays disjoint. + // 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 + // (gated below): prefill is compute-bound and gains nothing from the overlap. + const auto reach_backward = [](const ggml_tensor * start) { + std::unordered_set seen; + std::vector 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 < cgraph->n_nodes; ++join_idx) { + ggml_tensor * join_node = cgraph->nodes[join_idx]; + if (join_node->op != GGML_OP_ADD) { + continue; + } + + // Only overlap during decode (single token). Overlapping the shared expert only helps when + // the routed branch leaves the GPU underutilized for it to run alongside; that is the case in + // decode (batch 1, latency/occupancy-bound) but not in prefill, where the routed matmuls are + // large and already saturate the GPU, so the overlap adds contention without a speedup. + 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 reach_routed = reach_backward(routed_out); + const std::unordered_set 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 2, routed -> 1 + std::vector> nodes_per_branch(2); + for (int i = fork_idx + 1; i < join_idx; ++i) { + const ggml_tensor * n = cgraph->nodes[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; + } + + // the routed experts stay on the main stream and only the shared expert forks onto a single + // aux stream, joined at the add. Keeping the large routed branch on the main stream avoids + // migrating it and needs only one fork/join. + ggml_cuda_concurrent_event concurrent_event(1); + concurrent_event.join_node = join_node; + for (const ggml_tensor * n : nodes_per_branch[1]) { + concurrent_event.stream_mapping[n] = 1; + } + + const ggml_tensor * fork_node = cgraph->nodes[fork_idx]; + concurrent_event.original_order.reserve(join_idx - fork_idx - 1); + for (int i = fork_idx + 1; i < join_idx; ++i) { + concurrent_event.original_order.push_back(cgraph->nodes[i]); + } + + std::unordered_map & concurrent_events = cuda_ctx->stream_context().concurrent_events; + if (concurrent_events.find(fork_node) != concurrent_events.end()) { + continue; + } + concurrent_events.emplace(fork_node, std::move(concurrent_event)); + GGML_LOG_DEBUG("Adding shared-expert stream at node %s %p\n", fork_node->name, fork_node); + concurrent_node_ranges.emplace_back(fork_idx, join_idx); + + // the shared-expert nodes get a dedicated buffer (below), so the graph order is left intact + // and no interleaving is needed to keep the branch non-overlapping + concurrent_groups.push_back(nodes_per_branch[1]); + } + + // Place every concurrent branch (attention QKV and MoE shared-expert) in a dedicated buffer so + // its nodes never share an address with each other or with tensors read across the region (which + // ggml-alloc could otherwise recycle, corrupting concurrent reads). Layers run sequentially, so + // one buffer sized to the largest region is reused across all of them; within a region each node + // gets a distinct offset so the concurrent scratch stays disjoint. if (!concurrent_groups.empty()) { const size_t alignment = 128; From 6277fbefc26be592a565bb5817d266cc01327450 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Tue, 7 Jul 2026 06:44:02 -0700 Subject: [PATCH 4/6] ggml-cuda: drop dead concurrent-region reorder path The branch scratch buffer made the node interleaving obsolete, so the concurrent branches already run in their original graph order. The "restore original order" pass in ggml_cuda_graph_evaluate_and_capture is therefore a no-op, and the ggml_cuda_concurrent_event::original_order it consumed is now unused. Remove the reorder pass and original_order, keeping only the is_valid() check that clears the events when a region cannot be safely overlapped. Co-Authored-By: Claude Opus 4 (1M context) --- ggml/src/ggml-cuda/common.cuh | 5 --- ggml/src/ggml-cuda/ggml-cuda.cu | 64 +-------------------------------- 2 files changed, 1 insertion(+), 68 deletions(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index ffc01cc2ec79..630c831c3eca 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1244,10 +1244,6 @@ struct ggml_cuda_concurrent_event { int n_streams = 0; std::unordered_map stream_mapping; - // Original order of nodes in this concurrent region (before interleaving) - // Used to restore grouping for fusion within streams - std::vector 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; } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index b9d654c0d8ce..dbd765dc4321 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4287,58 +4287,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } } - if (should_launch_concurrent_events) { - // Restore original node order within each concurrent region to enable fusion within streams - - std::unordered_map node_to_idx; - node_to_idx.reserve(cgraph->n_nodes); - for (int i = 0; i < cgraph->n_nodes; ++i) { - node_to_idx[cgraph->nodes[i]] = i; - } - - for (auto & [fork_node, event] : stream_ctx.concurrent_events) { - // Find positions of all nodes from this event in the current graph - std::vector positions; - positions.reserve(event.original_order.size()); - - bool all_found = true; - for (const ggml_tensor * orig_node : event.original_order) { - auto it = node_to_idx.find(orig_node); - if (it != node_to_idx.end()) { - positions.push_back(it->second); - } else { - all_found = false; - break; - } - } - - if (!all_found || positions.size() != event.original_order.size()) { - continue; - } - - // Sort positions to get contiguous range - std::vector sorted_positions = positions; - std::sort(sorted_positions.begin(), sorted_positions.end()); - - bool is_contiguous = true; - for (size_t i = 1; i < sorted_positions.size(); ++i) { - if (sorted_positions[i] != sorted_positions[i-1] + 1) { - is_contiguous = false; - break; - } - } - - if (!is_contiguous) { - continue; - } - - // Restore original order at the sorted positions - int start_pos = sorted_positions[0]; - for (size_t i = 0; i < event.original_order.size(); ++i) { - cgraph->nodes[start_pos + i] = const_cast(event.original_order[i]); - } - } - } else { + if (!should_launch_concurrent_events) { stream_ctx.concurrent_events.clear(); } @@ -4758,13 +4707,6 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph continue; } - // Save the original order of nodes in this region before interleaving - // This is used later to restore grouping for fusion within streams - concurrent_event.original_order.reserve(total_branch_nodes); - for (int i = fork_node_idx + 1; i < join_node_idx; ++i) { - concurrent_event.original_order.push_back(cgraph->nodes[i]); - } - std::unordered_map & concurrent_events = cuda_ctx->stream_context().concurrent_events; GGML_ASSERT(concurrent_events.find(root_node) == concurrent_events.end()); concurrent_events.emplace(root_node, std::move(concurrent_event)); @@ -4885,10 +4827,6 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph } const ggml_tensor * fork_node = cgraph->nodes[fork_idx]; - concurrent_event.original_order.reserve(join_idx - fork_idx - 1); - for (int i = fork_idx + 1; i < join_idx; ++i) { - concurrent_event.original_order.push_back(cgraph->nodes[i]); - } std::unordered_map & concurrent_events = cuda_ctx->stream_context().concurrent_events; if (concurrent_events.find(fork_node) != concurrent_events.end()) { From 54e238e76cdb729ed5e1018751789e45e8e593ca Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Tue, 7 Jul 2026 08:22:23 -0700 Subject: [PATCH 5/6] ggml-cuda: extract concurrent-region detection into a pure pass Move the diamond-detection heuristics (attention QKV fan-out and the MoE shared-expert diamond) out of ggml_backend_cuda_graph_optimize into a header-only function ggml_cuda_detect_concurrent_regions() that operates only on the public ggml graph API and holds no CUDA state. The backend now turns the detected regions into ggml_cuda_concurrent_event objects and lays their branch tensors into the scratch buffer. Behavior is unchanged, and the detection heuristics can now be exercised in isolation. Co-Authored-By: Claude Opus 4 (1M context) --- ggml/src/ggml-cuda/concurrency.hpp | 331 +++++++++++++++++++++++++++++ ggml/src/ggml-cuda/ggml-cuda.cu | 304 ++------------------------ 2 files changed, 347 insertions(+), 288 deletions(-) create mode 100644 ggml/src/ggml-cuda/concurrency.hpp diff --git a/ggml/src/ggml-cuda/concurrency.hpp b/ggml/src/ggml-cuda/concurrency.hpp new file mode 100644 index 000000000000..a71fb5c89adf --- /dev/null +++ b/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 +#include +#include +#include +#include +#include + +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 stream_mapping; + + // branch nodes to place in the dedicated scratch buffer, in graph order + std::vector group; + + // true for the MoE shared-expert diamond, false for the attention QKV fan-out + bool shared_expert = false; +}; + +inline std::vector ggml_cuda_detect_concurrent_regions(ggml_cgraph * cgraph) { + std::vector 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 fan_out; + // reverse mapping of node -> index in the graph + std::unordered_map 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> concurrent_node_ranges; + // fork nodes already used as a region key (dedup across attention and MoE passes) + std::unordered_set 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> 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 & 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 & 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 seen; + std::vector 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 reach_routed = reach_backward(routed_out); + const std::unordered_set 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> 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; +} diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index dbd765dc4321..c0e018fc6c56 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4,6 +4,7 @@ #include "ggml-cuda/allreduce.cuh" #include "ggml-cuda/common.cuh" +#include "ggml-cuda/concurrency.hpp" #include "ggml-cuda/acc.cuh" #include "ggml-cuda/add-id.cuh" #include "ggml-cuda/arange.cuh" @@ -4540,305 +4541,32 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph return; } - // number of out-degrees for a particular node - std::unordered_map fan_out; - // reverse mapping of node to index in the cgraph - std::unordered_map node_indices; + // detect the concurrent regions (attention QKV fan-out and MoE shared-expert diamond) + std::vector regions = ggml_cuda_detect_concurrent_regions(cgraph); - const auto & is_noop = [](const ggml_tensor * node) -> bool { + 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 (uint32_t 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; - if (dst2 == src2) { - return true; - } - return false; - }; - - for (int node_idx = 0; node_idx < cgraph->n_nodes; node_idx++) { - const ggml_tensor * node = cgraph->nodes[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 = cgraph->nodes[node_idx]->src[src_idx]; - //TODO: check why nrows > 1 fails - if (node && !is_noop(node) && ggml_nrows(node) <= 1) { - fan_out[src] += 1; - } - } - } + std::unordered_map & concurrent_events = stream_context.concurrent_events; - // Target Q, K, V for concurrency - // this is a more general way to find nodes which can be candidates for concurrency (although it has not been tested for anything else): - // 1. find fan-out (fork) nodes where the same input is used at least N times (in QKV, it would be "attn-norm") - // 2. find the join node, where 2 or more of the outputs are required (in QKV, this would "KQ" or "flash-attn") - // 3. account for all branches from the fork to the join - // 4. To extend lifetimes of the tensors, we interleave the branches (see below for more details) - // 5. save the original cgraph and restore it in graph_compute, to enable fusion within streams - // See discussion: https://github.com/ggml-org/llama.cpp/pull/16991#issuecomment-3522620030 - - const int min_fan_out = 3; - const int max_fan_out = 3; - - // store {fork_idx, join_idx} - std::vector> 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 + // branch nodes of each region, placed in the dedicated scratch buffer (below) std::vector> concurrent_groups; + concurrent_groups.reserve(regions.size()); - 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]; - - // 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> nodes_per_branch; - for (int i = root_node_idx + 1; i < cgraph->n_nodes; ++i) { - const ggml_tensor * node = cgraph->nodes[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 & 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 < cgraph->n_nodes; ++i) { - const ggml_tensor * curr_node = cgraph->nodes[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 & branch_vec = nodes_per_branch[branch_idx]; - if (belongs_to_branch(curr_node, branch_vec)) { - //continue accumulating - 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) { - //Create ggml_cuda_concurrent_event - ggml_cuda_concurrent_event concurrent_event(nodes_per_branch.size()); - concurrent_event.join_node = join_node; - - for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) { - for (const ggml_tensor * n : nodes_per_branch[branch_idx]) { - concurrent_event.stream_mapping[n] = branch_idx + 1; - } - } - - int fork_node_idx = node_indices[root_node]; - int join_node_idx = node_indices[join_node]; - - int total_branch_nodes = 0; - for (std::vector branch_nodes : nodes_per_branch) { - total_branch_nodes += 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) { - GGML_LOG_DEBUG( - "Skipping %s because the number of nodes in the middle is not equal to the total number of " - "branch nodes %d != %d\n", - root_node->name, join_node_idx - fork_node_idx - 1, total_branch_nodes); - continue; - } + for (ggml_cuda_concurrent_region & region : regions) { + ggml_cuda_concurrent_event concurrent_event(region.n_streams); + concurrent_event.join_node = region.join_node; + concurrent_event.stream_mapping = std::move(region.stream_mapping); - std::unordered_map & concurrent_events = cuda_ctx->stream_context().concurrent_events; - GGML_ASSERT(concurrent_events.find(root_node) == concurrent_events.end()); - concurrent_events.emplace(root_node, std::move(concurrent_event)); - 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 group; - for (const auto & branch_nodes : nodes_per_branch) { - for (const ggml_tensor * n : branch_nodes) { - group.push_back(n); - } - } - concurrent_groups.push_back(std::move(group)); - } - } - } + GGML_ASSERT(concurrent_events.find(region.fork_node) == concurrent_events.end()); + concurrent_events.emplace(region.fork_node, std::move(concurrent_event)); - // 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 - // (gated below): prefill is compute-bound and gains nothing from the overlap. - const auto reach_backward = [](const ggml_tensor * start) { - std::unordered_set seen; - std::vector 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 < cgraph->n_nodes; ++join_idx) { - ggml_tensor * join_node = cgraph->nodes[join_idx]; - if (join_node->op != GGML_OP_ADD) { - continue; - } - - // Only overlap during decode (single token). Overlapping the shared expert only helps when - // the routed branch leaves the GPU underutilized for it to run alongside; that is the case in - // decode (batch 1, latency/occupancy-bound) but not in prefill, where the routed matmuls are - // large and already saturate the GPU, so the overlap adds contention without a speedup. - 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 reach_routed = reach_backward(routed_out); - const std::unordered_set 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 2, routed -> 1 - std::vector> nodes_per_branch(2); - for (int i = fork_idx + 1; i < join_idx; ++i) { - const ggml_tensor * n = cgraph->nodes[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; - } - - // the routed experts stay on the main stream and only the shared expert forks onto a single - // aux stream, joined at the add. Keeping the large routed branch on the main stream avoids - // migrating it and needs only one fork/join. - ggml_cuda_concurrent_event concurrent_event(1); - concurrent_event.join_node = join_node; - for (const ggml_tensor * n : nodes_per_branch[1]) { - concurrent_event.stream_mapping[n] = 1; - } - - const ggml_tensor * fork_node = cgraph->nodes[fork_idx]; - - std::unordered_map & concurrent_events = cuda_ctx->stream_context().concurrent_events; - if (concurrent_events.find(fork_node) != concurrent_events.end()) { - continue; - } - concurrent_events.emplace(fork_node, std::move(concurrent_event)); - GGML_LOG_DEBUG("Adding shared-expert stream at node %s %p\n", fork_node->name, fork_node); - concurrent_node_ranges.emplace_back(fork_idx, join_idx); + GGML_LOG_DEBUG("%s at node %s %p\n", region.shared_expert ? "Adding shared-expert stream" : "Adding stream", + region.fork_node->name, (const void *) region.fork_node); - // the shared-expert nodes get a dedicated buffer (below), so the graph order is left intact - // and no interleaving is needed to keep the branch non-overlapping - concurrent_groups.push_back(nodes_per_branch[1]); + concurrent_groups.push_back(std::move(region.group)); } // Place every concurrent branch (attention QKV and MoE shared-expert) in a dedicated buffer so From a86d31339cb5c61d0c3f5c68b5ec34cf9547a90e Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Tue, 7 Jul 2026 08:22:23 -0700 Subject: [PATCH 6/6] tests: add CUDA multi-stream concurrency tests Built only for CUDA/HIP/MUSA: - detection: drive ggml_cuda_detect_concurrent_regions() with synthetic diamond graphs and assert which regions are and are not detected - attention 3-way fan-out, MoE shared-expert diamond, and negatives (fan-out of 2 or 4, multi-row/prefill, wrong fork name, wrong MoE operand names). This also pins the tensor-name contract the detector relies on. - end-to-end: run the diamonds through ggml_backend_sched with GGML_CUDA_GRAPH_OPT=1 and check the output against an analytic reference; a log hook confirms both concurrency paths engaged. Skipped when no GPU is available. Co-Authored-By: Claude Opus 4 (1M context) --- tests/CMakeLists.txt | 6 + tests/test-cuda-concurrency.cpp | 396 ++++++++++++++++++++++++++++++++ 2 files changed, 402 insertions(+) create mode 100644 tests/test-cuda-concurrency.cpp diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 0dd1d7b162ae..ab87b3a4f920 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -242,6 +242,12 @@ endif() llama_build_and_test(test-gguf.cpp) llama_build_and_test(test-backend-ops.cpp) +# CUDA/HIP multi-stream concurrency: detection is a pure header pass, so it is unit-tested directly +if (GGML_CUDA OR GGML_HIP OR GGML_MUSA) + llama_build_and_test(test-cuda-concurrency.cpp) + target_include_directories(test-cuda-concurrency PRIVATE ${PROJECT_SOURCE_DIR}/ggml/src) +endif() + llama_build_and_test(test-model-load-cancel.cpp LABEL "model") llama_build_and_test(test-autorelease.cpp LABEL "model") llama_build_and_test(test-backend-sampler.cpp LABEL "model") diff --git a/tests/test-cuda-concurrency.cpp b/tests/test-cuda-concurrency.cpp new file mode 100644 index 000000000000..27b5a9a078b1 --- /dev/null +++ b/tests/test-cuda-concurrency.cpp @@ -0,0 +1,396 @@ +// Tests for the CUDA/HIP multi-stream concurrency support. +// +// ggml_cuda_detect_concurrent_regions() is a pure pass over a compute graph, so +// the detection heuristics are asserted directly against synthetic "diamond" +// graphs. The end-to-end test then runs the same diamonds on the backend and +// compares the result against a reference; it is skipped when no GPU is present. + +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" + +#include "ggml-cuda/concurrency.hpp" + +#include +#include +#include +#include +#include +#include + +static int g_failures = 0; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + g_failures++; \ + fprintf(stderr, " FAIL %s:%d: %s\n", __FILE__, __LINE__, #cond); \ + } \ + } while (0) + +struct graph_fixture { + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + + explicit graph_fixture(size_t mem = 16 * 1024 * 1024) { + ggml_init_params params = { mem, nullptr, /*no_alloc*/ true }; + ctx = ggml_init(params); + gf = ggml_new_graph(ctx); + } + ~graph_fixture() { ggml_free(ctx); } + + ggml_tensor * input_1d(int64_t n, const char * name) { + ggml_tensor * t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n); + ggml_set_name(t, name); + return t; + } + ggml_tensor * input_2d(int64_t n0, int64_t n1, const char * name) { + ggml_tensor * t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n0, n1); + ggml_set_name(t, name); + return t; + } + ggml_tensor * op(ggml_tensor * t, const char * name) { + ggml_set_name(t, name); + return t; + } + void expand(ggml_tensor * t) { ggml_build_forward_expand(gf, t); } +}; + +static const ggml_cuda_concurrent_region * find_by_fork(const std::vector & regions, + const char * fork_name) { + for (const auto & r : regions) { + if (r.fork_node && strstr(r.fork_node->name, fork_name)) { + return &r; + } + } + return nullptr; +} + +// ----------------------------------------------------------------------------- +// Attention QKV fan-out: fork "attn_norm" -> 3 branches -> join. +// Node order after expansion: [attn_norm, v, q, k, join]. +// ----------------------------------------------------------------------------- +static void test_attention_diamond() { + printf("test_attention_diamond\n"); + graph_fixture f; + + ggml_tensor * a = f.input_1d(128, "a"); + ggml_tensor * fork = f.op(ggml_rms_norm(f.ctx, a, 1e-5f), "attn_norm-0"); + ggml_tensor * q = f.op(ggml_scale(f.ctx, fork, 2.0f), "q-0"); + ggml_tensor * k = f.op(ggml_scale(f.ctx, fork, 3.0f), "k-0"); + ggml_tensor * v = f.op(ggml_scale(f.ctx, fork, 4.0f), "v-0"); + ggml_tensor * join = f.op(ggml_add(f.ctx, q, k), "kq-0"); + + f.expand(v); // pull v into the graph before the join + f.expand(join); + + auto regions = ggml_cuda_detect_concurrent_regions(f.gf); + + const ggml_cuda_concurrent_region * r = find_by_fork(regions, "attn_norm"); + CHECK(r != nullptr); + if (r) { + CHECK(r->shared_expert == false); + CHECK(r->n_streams == 3); + CHECK(r->join_node == join); + CHECK(r->stream_mapping.size() == 3); + CHECK(r->stream_mapping.count(q) == 1); + CHECK(r->stream_mapping.count(k) == 1); + CHECK(r->stream_mapping.count(v) == 1); + // the three branches must land on three distinct aux streams + CHECK(r->stream_mapping.at(q) != r->stream_mapping.at(k)); + CHECK(r->stream_mapping.at(q) != r->stream_mapping.at(v)); + CHECK(r->stream_mapping.at(k) != r->stream_mapping.at(v)); + for (const auto & [node, stream] : r->stream_mapping) { + CHECK(stream >= 1 && stream <= 3); + (void) node; + } + } +} + +// ----------------------------------------------------------------------------- +// MoE shared-expert diamond: fork "ffn_norm" -> {routed, shared} -> add. +// Node order after expansion: [ffn_norm, ffn_moe_out, shexp_gate, ffn_shexp, add]. +// ----------------------------------------------------------------------------- +static void test_moe_shexp_diamond() { + printf("test_moe_shexp_diamond\n"); + graph_fixture f; + + ggml_tensor * x = f.input_1d(128, "ffn_inp"); + ggml_tensor * fork = f.op(ggml_rms_norm(f.ctx, x, 1e-5f), "ffn_norm-3"); + ggml_tensor * routed = f.op(ggml_scale(f.ctx, fork, 2.0f), "ffn_moe_out-3"); + ggml_tensor * gate = f.op(ggml_scale(f.ctx, fork, 3.0f), "ffn_shexp_gate-3"); + ggml_tensor * shexp = f.op(ggml_scale(f.ctx, gate, 1.5f), "ffn_shexp-3"); + ggml_tensor * join = f.op(ggml_add(f.ctx, routed, shexp), "ffn_out-3"); + + f.expand(join); + + auto regions = ggml_cuda_detect_concurrent_regions(f.gf); + + const ggml_cuda_concurrent_region * r = find_by_fork(regions, "ffn_norm"); + CHECK(r != nullptr); + if (r) { + CHECK(r->shared_expert == true); + CHECK(r->n_streams == 1); + CHECK(r->join_node == join); + CHECK(r->fork_node == fork); + // only the shared-expert branch is mapped; routed stays on the main stream + CHECK(r->stream_mapping.count(shexp) == 1); + CHECK(r->stream_mapping.count(gate) == 1); + CHECK(r->stream_mapping.count(routed) == 0); + CHECK(r->stream_mapping.at(shexp) == 1); + } +} + +// two-way fan-out must NOT be detected (min_fan_out is 3) +static void test_negative_fanout_two() { + printf("test_negative_fanout_two\n"); + graph_fixture f; + + ggml_tensor * a = f.input_1d(128, "a"); + ggml_tensor * fork = f.op(ggml_rms_norm(f.ctx, a, 1e-5f), "attn_norm-0"); + ggml_tensor * q = f.op(ggml_scale(f.ctx, fork, 2.0f), "q-0"); + ggml_tensor * k = f.op(ggml_scale(f.ctx, fork, 3.0f), "k-0"); + ggml_tensor * join = f.op(ggml_add(f.ctx, q, k), "kq-0"); + f.expand(join); + + auto regions = ggml_cuda_detect_concurrent_regions(f.gf); + CHECK(find_by_fork(regions, "attn_norm") == nullptr); +} + +// four-way fan-out must NOT be detected (max_fan_out is 3) +static void test_negative_fanout_four() { + printf("test_negative_fanout_four\n"); + graph_fixture f; + + ggml_tensor * a = f.input_1d(128, "a"); + ggml_tensor * fork = f.op(ggml_rms_norm(f.ctx, a, 1e-5f), "attn_norm-0"); + ggml_tensor * q = f.op(ggml_scale(f.ctx, fork, 2.0f), "q-0"); + ggml_tensor * k = f.op(ggml_scale(f.ctx, fork, 3.0f), "k-0"); + ggml_tensor * v = f.op(ggml_scale(f.ctx, fork, 4.0f), "v-0"); + ggml_tensor * w = f.op(ggml_scale(f.ctx, fork, 5.0f), "w-0"); + ggml_tensor * j1 = f.op(ggml_add(f.ctx, q, k), "j1-0"); + ggml_tensor * join = f.op(ggml_add(f.ctx, j1, v), "kq-0"); + f.expand(w); + f.expand(join); + + auto regions = ggml_cuda_detect_concurrent_regions(f.gf); + CHECK(find_by_fork(regions, "attn_norm") == nullptr); +} + +// prefill (nrows > 1) must NOT be detected: fan-out only counts single-row consumers +static void test_negative_multirow() { + printf("test_negative_multirow\n"); + graph_fixture f; + + ggml_tensor * a = f.input_2d(128, 8, "a"); + ggml_tensor * fork = f.op(ggml_rms_norm(f.ctx, a, 1e-5f), "attn_norm-0"); + ggml_tensor * q = f.op(ggml_scale(f.ctx, fork, 2.0f), "q-0"); + ggml_tensor * k = f.op(ggml_scale(f.ctx, fork, 3.0f), "k-0"); + ggml_tensor * v = f.op(ggml_scale(f.ctx, fork, 4.0f), "v-0"); + ggml_tensor * join = f.op(ggml_add(f.ctx, q, k), "kq-0"); + f.expand(v); + f.expand(join); + + auto regions = ggml_cuda_detect_concurrent_regions(f.gf); + CHECK(find_by_fork(regions, "attn_norm") == nullptr); +} + +// right shape, wrong fork name -> attention pass must ignore it +static void test_negative_wrong_name() { + printf("test_negative_wrong_name\n"); + graph_fixture f; + + ggml_tensor * a = f.input_1d(128, "a"); + ggml_tensor * fork = f.op(ggml_rms_norm(f.ctx, a, 1e-5f), "some_norm-0"); + ggml_tensor * q = f.op(ggml_scale(f.ctx, fork, 2.0f), "q-0"); + ggml_tensor * k = f.op(ggml_scale(f.ctx, fork, 3.0f), "k-0"); + ggml_tensor * v = f.op(ggml_scale(f.ctx, fork, 4.0f), "v-0"); + ggml_tensor * join = f.op(ggml_add(f.ctx, q, k), "kq-0"); + f.expand(v); + f.expand(join); + + auto regions = ggml_cuda_detect_concurrent_regions(f.gf); + CHECK(regions.empty()); +} + +// MoE add whose operands don't carry the expected names -> not detected +static void test_negative_moe_wrong_names() { + printf("test_negative_moe_wrong_names\n"); + graph_fixture f; + + ggml_tensor * x = f.input_1d(128, "ffn_inp"); + ggml_tensor * fork = f.op(ggml_rms_norm(f.ctx, x, 1e-5f), "ffn_norm-3"); + ggml_tensor * routed = f.op(ggml_scale(f.ctx, fork, 2.0f), "ffn_moe_out-3"); + ggml_tensor * other = f.op(ggml_scale(f.ctx, fork, 3.0f), "ffn_other-3"); + ggml_tensor * join = f.op(ggml_add(f.ctx, routed, other), "ffn_out-3"); + f.expand(join); + + auto regions = ggml_cuda_detect_concurrent_regions(f.gf); + CHECK(regions.empty()); +} + +// ----------------------------------------------------------------------------- +// End-to-end: run the diamonds on the backend via ggml_backend_sched (which is +// what invokes graph_optimize) with GGML_CUDA_GRAPH_OPT=1, and check the output +// against an analytic reference. graph_optimize logs when it forks a region, so a +// log hook is used to confirm the concurrency path actually engaged. +// ----------------------------------------------------------------------------- +static bool g_attn_engaged = false; +static bool g_shexp_engaged = false; + +static void concurrency_log_hook(ggml_log_level level, const char * text, void * /*user*/) { + (void) level; + if (!text) { + return; + } + if (strstr(text, "Adding shared-expert stream")) { + g_shexp_engaged = true; + } else if (strstr(text, "Adding stream")) { + g_attn_engaged = true; + } +} + +static void rms_norm_ref(const std::vector & x, std::vector & out, float eps) { + double sum = 0.0; + for (float v : x) { + sum += (double) v * (double) v; + } + const double scale = 1.0 / std::sqrt(sum / (double) x.size() + (double) eps); + out.resize(x.size()); + for (size_t i = 0; i < x.size(); ++i) { + out[i] = (float) ((double) x[i] * scale); + } +} + +static void test_execution_correctness() { + printf("test_execution_correctness\n"); + + ggml_backend_load_all(); + + // accept a discrete or integrated GPU (Strix Halo reports IGPU) + ggml_backend_dev_t gpu_dev = nullptr; + for (size_t i = 0; i < ggml_backend_dev_count() && !gpu_dev; ++i) { + ggml_backend_dev_t d = ggml_backend_dev_get(i); + const auto type = ggml_backend_dev_type(d); + if (type == GGML_BACKEND_DEVICE_TYPE_GPU || type == GGML_BACKEND_DEVICE_TYPE_IGPU) { + gpu_dev = d; + } + } + ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); + if (!gpu_dev || !cpu_dev) { + printf(" SKIP: no GPU/CPU device available\n"); + return; + } + + ggml_backend_t gpu = ggml_backend_dev_init(gpu_dev, nullptr); + ggml_backend_t cpu = ggml_backend_dev_init(cpu_dev, nullptr); + + const int64_t n = 256; + const float eps = 1e-5f; + + ggml_init_params params = { ggml_tensor_overhead() * 64 + ggml_graph_overhead(), nullptr, /*no_alloc*/ true }; + ggml_context * ctx = ggml_init(params); + + // attention-like diamond: join_attn = q + k = 5 * rms_norm(a) + ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n); + ggml_set_name(a, "a"); ggml_set_input(a); + ggml_tensor * fork_a = ggml_rms_norm(ctx, a, eps); ggml_set_name(fork_a, "attn_norm-0"); + ggml_tensor * q = ggml_scale(ctx, fork_a, 2.0f); ggml_set_name(q, "q-0"); + ggml_tensor * k = ggml_scale(ctx, fork_a, 3.0f); ggml_set_name(k, "k-0"); + ggml_tensor * v = ggml_scale(ctx, fork_a, 4.0f); ggml_set_name(v, "v-0"); + ggml_tensor * join_attn = ggml_add(ctx, q, k); ggml_set_name(join_attn, "kq-0"); + ggml_set_output(join_attn); + + // MoE-like diamond: join_moe = routed + shexp = (2 + 4.5) * rms_norm(b) + ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n); + ggml_set_name(b, "ffn_inp"); ggml_set_input(b); + ggml_tensor * fork_b = ggml_rms_norm(ctx, b, eps); ggml_set_name(fork_b, "ffn_norm-3"); + ggml_tensor * routed = ggml_scale(ctx, fork_b, 2.0f); ggml_set_name(routed, "ffn_moe_out-3"); + ggml_tensor * gate = ggml_scale(ctx, fork_b, 3.0f); ggml_set_name(gate, "ffn_shexp_gate-3"); + ggml_tensor * shexp = ggml_scale(ctx, gate, 1.5f); ggml_set_name(shexp, "ffn_shexp-3"); + ggml_tensor * join_moe = ggml_add(ctx, routed, shexp); ggml_set_name(join_moe, "ffn_out-3"); + ggml_set_output(join_moe); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, v); // keep v in the region, before the attn join + ggml_build_forward_expand(gf, join_attn); + ggml_build_forward_expand(gf, join_moe); + + // allocate the whole graph on the GPU so the scheduler keeps every op there (and + // thus runs graph_optimize on a single GPU split, the way decode does) + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, gpu); + CHECK(buf != nullptr); + + std::vector adata(n), bdata(n); + for (int64_t i = 0; i < n; ++i) { + adata[i] = 0.5f + 0.01f * (float) (i % 17); + bdata[i] = -0.3f + 0.02f * (float) (i % 13); + } + ggml_backend_tensor_set(a, adata.data(), 0, n * sizeof(float)); + ggml_backend_tensor_set(b, bdata.data(), 0, n * sizeof(float)); + + ggml_backend_t backends[] = { gpu, cpu }; + ggml_backend_buffer_type_t bufts[] = { ggml_backend_dev_buffer_type(gpu_dev), + ggml_backend_dev_buffer_type(cpu_dev) }; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, bufts, 2, 2048, false, false); + + // a couple of iterations to pass CUDA-graph warmup and exercise capture + replay + for (int iter = 0; iter < 3; ++iter) { + ggml_backend_sched_reset(sched); + CHECK(ggml_backend_sched_graph_compute(sched, gf) == GGML_STATUS_SUCCESS); + } + + std::vector ref_a, ref_b; + rms_norm_ref(adata, ref_a, eps); + rms_norm_ref(bdata, ref_b, eps); + + std::vector out_attn(n), out_moe(n); + ggml_backend_tensor_get(join_attn, out_attn.data(), 0, n * sizeof(float)); + ggml_backend_tensor_get(join_moe, out_moe.data(), 0, n * sizeof(float)); + + const float tol = 1e-3f; + int bad = 0; + for (int64_t i = 0; i < n; ++i) { + if (std::fabs(out_attn[i] - 5.0f * ref_a[i]) > tol) { bad++; } + if (std::fabs(out_moe[i] - 6.5f * ref_b[i]) > tol) { bad++; } + } + CHECK(bad == 0); + + if (g_attn_engaged) { + printf(" attention concurrency engaged\n"); + } else { + printf(" NOTE: attention concurrency did not engage (CUDA graphs off?) - correctness still checked\n"); + } + if (g_shexp_engaged) { + printf(" shared-expert concurrency engaged\n"); + } else { + printf(" NOTE: shared-expert concurrency did not engage - correctness still checked\n"); + } + + ggml_backend_sched_free(sched); + ggml_backend_buffer_free(buf); + ggml_free(ctx); + ggml_backend_free(gpu); + ggml_backend_free(cpu); +} + +int main() { + test_attention_diamond(); + test_moe_shexp_diamond(); + test_negative_fanout_two(); + test_negative_fanout_four(); + test_negative_multirow(); + test_negative_wrong_name(); + test_negative_moe_wrong_names(); + + setenv("GGML_CUDA_GRAPH_OPT", "1", 1); + ggml_log_set(concurrency_log_hook, nullptr); + test_execution_correctness(); + + if (g_failures == 0) { + printf("all tests passed\n"); + return 0; + } + fprintf(stderr, "%d checks failed\n", g_failures); + return 1; +}