Commit 280141e6d for llama.cpp
commit 280141e6d4769b1e4424de999437cc72ffac33d6
Author: Aman Gupta <amangupta052@gmail.com>
Date: Sat Oct 24 00:26:16 2026 +0800
CUDA: chunk graphs to amortize graph launch (#30804)
diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh
index 049100fe3..37d0d104c 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<cudaGraphExec_t> instances;
size_t num_nodes = 0;
std::vector<cudaGraphNode_t> 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 e43807392..ff9e9a7ea 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<std::mutex> 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 (cuda_graph_update_required) { // Update graph executable
- ggml_cuda_graph_update_executable(cuda_ctx, graph_key);
+ 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()));
}
- // 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();
}
}
}