diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 049100fe38..37d0d104cd 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1264,15 +1264,12 @@ struct ggml_tensor_extra_gpu { struct ggml_cuda_graph { #ifdef USE_CUDA_GRAPH ~ggml_cuda_graph() { - if (instance != nullptr) { + for (cudaGraphExec_t instance : instances) { CUDA_CHECK(cudaGraphExecDestroy(instance)); } - if (graph != nullptr) { - CUDA_CHECK(cudaGraphDestroy(graph)); - } } - cudaGraph_t graph = nullptr; - cudaGraphExec_t instance = nullptr; + // the ggml graph is captured as a series of CUDA graphs, so the GPU can run one chunk while the next is enqueued + std::vector instances; size_t num_nodes = 0; std::vector nodes; bool disable_due_to_gpu_arch = false; @@ -1501,7 +1498,7 @@ struct ggml_backend_cuda_context { // Check if any CUDA graph has an instance for this context bool any_cuda_graph_has_instance() const { for (const auto & [key, graph] : cuda_graphs) { - if (graph && graph->instance != nullptr) { + if (graph && !graph->instances.empty()) { return true; } } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index e43807392b..ff9e9a7eae 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2803,16 +2803,14 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx return res; } -static void ggml_cuda_graph_update_executable(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - +static void ggml_cuda_graph_update_executable(cudaGraphExec_t & instance, cudaGraph_t cuda_graph) { #if CUDART_VERSION >= 12000 cudaGraphExecUpdateResultInfo result_info; - cudaError_t stat = cudaGraphExecUpdate(graph->instance, graph->graph, &result_info); + cudaError_t stat = cudaGraphExecUpdate(instance, cuda_graph, &result_info); #else cudaGraphNode_t errorNode; cudaGraphExecUpdateResult result_info; - cudaError_t stat = cudaGraphExecUpdate(graph->instance, graph->graph, &errorNode, &result_info); + cudaError_t stat = cudaGraphExecUpdate(instance, cuda_graph, &errorNode, &result_info); #endif // CUDART_VERSION >= 12000 if (stat == cudaErrorGraphExecUpdateFailure) { @@ -2823,13 +2821,39 @@ static void ggml_cuda_graph_update_executable(ggml_backend_cuda_context * cuda_c // The pre-existing graph exec cannot be updated due to violated constraints // so instead clear error and re-instantiate (void)cudaGetLastError(); - CUDA_CHECK(cudaGraphExecDestroy(graph->instance)); - graph->instance = nullptr; - CUDA_CHECK(cudaGraphInstantiate(&graph->instance, graph->graph, NULL, NULL, 0)); + CUDA_CHECK(cudaGraphExecDestroy(instance)); + instance = nullptr; + CUDA_CHECK(cudaGraphInstantiate(&instance, cuda_graph, nullptr, nullptr, 0)); } else { GGML_ASSERT(stat == cudaSuccess); } } + +// number of ggml nodes per CUDA graph chunk, 0 captures the whole ggml graph as one CUDA graph +static int ggml_cuda_graph_chunk_size() { + static const int chunk_size = [] { + const char * env = getenv("GGML_CUDA_GRAPH_CHUNK_SIZE"); + return env != nullptr ? atoi(env) : 512; + }(); + return chunk_size; +} + +// end the capture of chunk `idx`, update or create its executable and launch it +static void ggml_cuda_graph_launch_chunk(ggml_backend_cuda_context * cuda_ctx, ggml_cuda_graph * graph, size_t idx) { + cudaGraph_t cuda_graph = nullptr; + CUDA_CHECK(cudaStreamEndCapture(cuda_ctx->stream(), &cuda_graph)); + + if (idx < graph->instances.size()) { + ggml_cuda_graph_update_executable(graph->instances[idx], cuda_graph); + } else { + GGML_ASSERT(idx == graph->instances.size()); + graph->instances.emplace_back(nullptr); + CUDA_CHECK(cudaGraphInstantiate(&graph->instances[idx], cuda_graph, nullptr, nullptr, 0)); + } + CUDA_CHECK(cudaGraphDestroy(cuda_graph)); + + CUDA_CHECK(cudaGraphLaunch(graph->instances[idx], cuda_ctx->stream())); +} #endif // USE_CUDA_GRAPH static bool ggml_cuda_should_fuse_rope_set_rows(const ggml_tensor * rope, @@ -4583,6 +4607,13 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud ggml_cuda_concurrent_event * concurrent_event = nullptr; bool should_launch_concurrent_events = false; +#ifdef USE_CUDA_GRAPH + ggml_cuda_graph * graph = use_cuda_graph ? cuda_ctx->cuda_graph(graph_key) : nullptr; + const int chunk_size = ggml_cuda_graph_chunk_size(); + int chunk_end = chunk_size; + size_t n_chunks = 0; +#endif // USE_CUDA_GRAPH + const auto try_launch_concurrent_event = [&](const ggml_tensor * node) { if (stream_ctx.concurrent_events.find(node) != stream_ctx.concurrent_events.end()) { concurrent_event = &stream_ctx.concurrent_events[node]; @@ -4672,6 +4703,15 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud for (int i = 0; i < cgraph->n_nodes; i++) { ggml_tensor * node = cgraph->nodes[i]; +#ifdef USE_CUDA_GRAPH + // fused nodes are skipped together, so a cut never splits a fusion + // a concurrent region must join before its capture can end + if (use_cuda_graph && chunk_size > 0 && i >= chunk_end && !is_concurrent_event_active) { + ggml_cuda_graph_launch_chunk(cuda_ctx, graph, n_chunks++); + CUDA_CHECK(cudaStreamBeginCapture(cuda_ctx->stream(), cudaStreamCaptureModeRelaxed)); + chunk_end = i + chunk_size; + } +#endif // USE_CUDA_GRAPH if (is_concurrent_event_active) { GGML_ASSERT(concurrent_event); @@ -4754,14 +4794,14 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } #ifdef USE_CUDA_GRAPH - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); if (use_cuda_graph && cuda_graph_update_required) { // End CUDA graph capture - if (graph->graph != nullptr) { - CUDA_CHECK(cudaGraphDestroy(graph->graph)); - graph->graph = nullptr; - } + ggml_cuda_graph_launch_chunk(cuda_ctx, graph, n_chunks++); - CUDA_CHECK(cudaStreamEndCapture(cuda_ctx->stream(), &graph->graph)); + // drop the chunks left over from a longer graph + for (size_t i = n_chunks; i < graph->instances.size(); ++i) { + CUDA_CHECK(cudaGraphExecDestroy(graph->instances[i])); + } + graph->instances.resize(n_chunks); graph_evaluated_or_captured = true; // CUDA graph has been captured std::lock_guard lock(ggml_cuda_lock); @@ -4773,16 +4813,11 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } } - if (use_cuda_graph) { - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - if (graph->instance == nullptr) { // Create executable graph from captured graph. - CUDA_CHECK(cudaGraphInstantiate(&graph->instance, graph->graph, NULL, NULL, 0)); + if (use_cuda_graph && !cuda_graph_update_required) { + // the chunks were launched during capture, otherwise launch them here + for (cudaGraphExec_t instance : graph->instances) { + CUDA_CHECK(cudaGraphLaunch(instance, cuda_ctx->stream())); } - if (cuda_graph_update_required) { // Update graph executable - ggml_cuda_graph_update_executable(cuda_ctx, graph_key); - } - // Launch graph - CUDA_CHECK(cudaGraphLaunch(graph->instance, cuda_ctx->stream())); #else GGML_UNUSED(graph_key); graph_evaluated_or_captured = true; @@ -4794,7 +4829,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - if (graph->graph == nullptr) { + if (graph->instances.empty()) { if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_VOLTA) { if (!graph->disable_due_to_gpu_arch && ggml_cuda_graph_debug_enabled()) { GGML_LOG_DEBUG("%s: disabling CUDA graphs due to GPU architecture\n", __func__); @@ -4848,7 +4883,7 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, } } else { use_cuda_graph = true; - cuda_graph_update_required = graph->instance == nullptr; + cuda_graph_update_required = graph->instances.empty(); } } }