diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index e6e50e041195..ffc01cc2ec79 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; @@ -1457,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) { @@ -1472,12 +1479,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..1b760476fe00 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"); @@ -614,10 +615,15 @@ 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])); + } } } + if (concurrent_scratch != nullptr) { + ggml_backend_buffer_free(concurrent_scratch); + } } @@ -1925,7 +1931,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 +2007,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)) { @@ -4349,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) { @@ -4359,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); } } @@ -4549,6 +4563,132 @@ static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_ev } } +// 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. +static void ggml_cuda_detect_shared_expert_concurrency( + ggml_cgraph * cgraph, + ggml_backend_cuda_context * cuda_ctx, + const std::unordered_map & node_indices, + std::vector> & concurrent_node_ranges, + std::vector> & concurrent_groups) { + 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 in the mat-vec regime (up to MMVQ_MAX_BATCH_SIZE tokens). The shared-expert + // overlap only helps when the routed branch leaves the GPU underutilized for it to run + // alongside; that holds while the matmuls stay in the memory/occupancy-bound mat-vec path, + // but not once the token count grows past it and the routed matmuls saturate the GPU + // (prefill), where the overlap only adds contention. + if (ggml_nrows(join_node) > MMVQ_MAX_BATCH_SIZE) { + 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]); + } +} + static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; @@ -4561,11 +4701,17 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph GGML_UNUSED(cgraph); #endif - static bool enable_graph_optimization = [] { - const char * env = getenv("GGML_CUDA_GRAPH_OPT"); - return env != nullptr && atoi(env) == 1; + // GGML_CUDA_GRAPH_OPT: "1"/"0" force the optimization on/off; when unset it defaults on for + // RDNA3.5, the only architecture where the decode-time overlap has been tuned, and off elsewhere. + static const int graph_opt_env = [] { + const char * env = getenv("GGML_CUDA_GRAPH_OPT"); + return env == nullptr ? -1 : atoi(env); }(); + const bool enable_graph_optimization = graph_opt_env >= 0 ? + graph_opt_env == 1 : + GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc); + if (!enable_graph_optimization) { return; } @@ -4633,6 +4779,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]; @@ -4724,10 +4875,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(); @@ -4756,37 +4903,66 @@ 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); + ggml_cuda_detect_shared_expert_concurrency(cgraph, cuda_ctx, node_indices, concurrent_node_ranges, concurrent_groups); - if (branch_nodes.empty()) { - current_branch_idx = (current_branch_idx + 1) % n_branches; - continue; - } + // 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; - cgraph->nodes[current_node_idx] = const_cast(branch_nodes.front()); - current_node_idx++; - branch_nodes.erase(branch_nodes.begin()); + 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; + }; - // 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()); - } + size_t needed = 0; + for (const auto & group : concurrent_groups) { + needed = std::max(needed, group_footprint(group)); + } + + 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; + } - current_branch_idx = (current_branch_idx + 1) % n_branches; + 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; + } + ggml_tensor * n = const_cast(cn); + n->data = base + off; + n->buffer = cuda_ctx->concurrent_scratch; + off += GGML_PAD(ggml_nbytes(n), alignment); } } }