From b966e226e3caf5c57255331ba78b420b7b170b25 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Sun, 5 Jul 2026 00:04:22 -0700 Subject: [PATCH 1/3] 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 | 6 ++++-- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 290dc4aff259..f7732ffd30e6 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 413ceddb7ce6..426de91dbb50 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -648,8 +648,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])); + } } } } From fd0b5b40aa2e3043d02d58f6b02340c6f15a2911 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Sun, 5 Jul 2026 00:31:45 -0700 Subject: [PATCH 2/3] 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 f7732ffd30e6..ac3da98627bc 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 426de91dbb50..beeddadd4e1a 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -654,6 +654,9 @@ ggml_backend_cuda_context::~ggml_backend_cuda_context() { } } } + if (concurrent_scratch != nullptr) { + ggml_backend_buffer_free(concurrent_scratch); + } } @@ -4188,6 +4191,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]; @@ -4279,10 +4287,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(); @@ -4311,37 +4315,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 7fc317fb3cef9021684c29741daaaf3738b8f5a5 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Sun, 5 Jul 2026 00:32:13 -0700 Subject: [PATCH 3/3] 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. Following vLLM's shared-expert model, 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 beeddadd4e1a..ca424c75ad2b 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -89,6 +89,7 @@ #include #include #include +#include #include static_assert(sizeof(half) == sizeof(ggml_fp16_t), "wrong fp16 size"); @@ -3900,8 +3901,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 (the vLLM shared-expert model) + 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) { @@ -3910,7 +3914,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); } } @@ -4329,11 +4334,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; + } + + // vLLM shared-expert model: 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;