perf(cuda): cap the CUDA graphs a context keeps by their nodes, not their count

TORAD.md

@@ -106,6 +106,7 @@ which pins a commit of this branch as a submodule.
 | `78e0fadf7` | `topk_radix` also takes a row of up to 1,024 columns at any k: with CUB's top-k available such a row had no path but CUB one row after another, four launches a row (a DSA indexer's ubatch under 4,096 cells: 2,048 launches a layer). RTX 5080: 3.3-3.9 us against 8.2-12.7 (1 row) and 129-171 (16 rows). | `GGML_CUDA_TOPK_RADIX_LEGACY` |
 | `640a12c41` | A SwiGLU limit (GLM-5.3's routed experts, shared experts and dense FFN: the gate clamped to `[-inf, 10]`, the up projection to `[-10, 10]`, then the GLU) fuses into the quantized mat-vec's GLU epilogue: `ggml_cuda_can_fuse` takes `{MUL_MAT(_ID), CLAMP, MUL_MAT(_ID), CLAMP, GLU}` when the clamps are exactly a limit and the GLU a SwiGLU, and `mul_mat_vec_q` applies it as one float, where the two nodes had kept the gate/up fusion from matching (five launches for one). In place (GLM-5.3-Flash's routed FFN, 8 layers of 32 experts, RTX 5080): 807 against 861 us at 1 token; a verify's 3 tokens are unchanged, their experts unfused. | `GGML_CUDA_GLU_LIMIT_FUSE_LEGACY` |
 | `b8d44d4b8` | Routed experts at 1-8 tokens stream through `mmvq-moe.cu`: one block an SM, a producer warp listing the distinct experts once (every token/slot pair that routes to one reads it once) and streaming tiles of their rows through a ring of shared-memory slots by `cp.async.bulk`, two teams of eight warps on the dot products; an IQ3_XXS warp frees its slot before the math, and the last tiles go by tickets on the stream's tile counter. It takes IQ3_XXS, and Q8_0 past 1 token: in place (GLM-5.3-Flash's routed FFN on 32 experts, CUDA graphs and PDL, RTX 5080) IQ3_XXS 815 against 818 us for 8 layers at 1 token and 1,823 against 1,953 at 3, Q8_0 2,265 against 2,382 for 4 layers at 3; IQ4_XS measured slower and keeps `mul_mat_vec_q`. Its steady state runs at DRAM's ceiling (~913 GB/s); a launch's first q8_1 vector, loaded from global memory behind the stream, costs 3-6 us. | `GGML_CUDA_MMVQ_MOE_LEGACY` |
+| `83f230ebb` | The CUDA graphs a context keeps are capped by their nodes: `GGML_CUDA_GRAPH_NODES` (default 32768) bounds the nodes their executables hold, evicting the least recently used, and `GGML_CUDA_GRAPH_MAX` bounds their count only when set. An executable holds a few KiB of device memory a node (1.9-3.0 on an RTX 5080), from a driver pool that never shrinks and stays at the most nodes held at once: 14, 44 and 88 MiB at 4,096, 16,384 and 32,768 nodes, under 3,500 evictions of mixed sizes. d0f8bae41's cap of 8 graphs broke `-sm tensor`: it computes a graph per all-reduce step on each device, two a layer every token, and each was evicted before its turn came round again, captured anew every token (qwen3-0.6b Q4_0 on a 5080 + 5070 Ti, 257.0 tok/s against 558.0 uncapped). `test-cuda-graph-cap-nodes`: twelve graphs of 8 nodes computed in turn under a cap of 100 nodes, each captured once and then replayed (a cap of 8 graphs captures 4 anew every time round). | `GGML_CUDA_GRAPH_NODES=0 GGML_CUDA_GRAPH_MAX=8` |
 
 Every switch in the last column is read once per process and parses as an integer: a `*_LEGACY` switch set to `0`
 is the same as unset (the change stays on), and `=0` turns off `GGML_CUDA_LORA_RANK1_FUSE` and

ggml/src/ggml-cuda/common.cuh

@@ -1300,7 +1300,7 @@ struct ggml_cuda_graph {
     }
     cudaGraph_t graph = nullptr;
     cudaGraphExec_t instance = nullptr;
-    size_t num_nodes = 0;
+    size_t num_nodes = 0; // the captured graph's: its executable holds device memory for each (cuda_graph_make_room)
     std::vector<cudaGraphNode_t> nodes;
     bool disable_due_to_gpu_arch = false;
     bool warmup_complete = false;
@@ -1746,38 +1746,60 @@ struct ggml_backend_cuda_context {
         return it->second.get();
     }
 
-    // The most CUDA graph executables a context keeps (GGML_CUDA_GRAPH_MAX, 0 for no cap). Each holds device memory
-    // beside the buffers the scheduler sizes, and the graphs a context computes follow its traffic (one per graph shape,
-    // and per graph slot for the same shape), so without a count they would be bounded only by the 10 s eviction above:
-    // a new one evicts the least recently used.
+    // The CUDA graph executables a context keeps. Each holds device memory beside the buffers the scheduler sizes, a few
+    // KiB for each node of its graph (2.75 on an RTX 5080), from a pool of the driver's that never shrinks: destroying
+    // one returns nothing to the device, a later one reuses it, and the pool stays at the most nodes held at once. The
+    // graphs a context computes follow its traffic (one per graph shape, and per graph slot for the same shape), bounded
+    // otherwise only by the 10 s eviction above, so a new one evicts the least recently used others until the nodes held
+    // fit under GGML_CUDA_GRAPH_NODES (default 32768, 88 MiB on that card; 0 for no cap) and the executables under
+    // GGML_CUDA_GRAPH_MAX (default 0, no cap). The memory follows the nodes, not the count: a -sm tensor split computes a
+    // small graph for each all-reduce step on each device, hundreds a token, and a cap on the count below that evicts
+    // every one of them before its turn comes round again, to be captured anew every token.
     static size_t cuda_graphs_max() {
         static const size_t n = [] {
             const char * env = getenv("GGML_CUDA_GRAPH_MAX");
-            return env != nullptr ? (size_t) std::max(0, atoi(env)) : (size_t) 8;
+            return env != nullptr ? (size_t) std::max(0, atoi(env)) : (size_t) 0;
         }();
         return n;
     }
-    size_t cuda_graphs_held_max = 0; // the most executables held at once here
-    size_t cuda_graph_bytes_max = 0; // the most device memory one took to instantiate here (cudaMemGetInfo around it)
+    static size_t cuda_graph_nodes_max() {
+        static const size_t n = [] {
+            const char * env = getenv("GGML_CUDA_GRAPH_NODES");
+            return env != nullptr ? (size_t) std::max(0, atoi(env)) : (size_t) 32768;
+        }();
+        return n;
+    }
+    size_t cuda_graphs_held_max      = 0; // the most executables held at once here
+    size_t cuda_graph_nodes_held_max = 0; // the most nodes they held at once here
+    size_t cuda_graph_bytes_max      = 0; // the most device memory one took to instantiate here (cudaMemGetInfo around it)
+    size_t cuda_graph_bytes          = 0; // and all of them
+
+    struct cuda_graphs_held {
+        size_t graphs;
+        size_t nodes;
+    };
 
-    // before `graph` gets an executable: evict the least recently used others until it fits under the cap; returns how
-    // many are held with it
-    size_t cuda_graph_make_room(const ggml_cuda_graph * graph) {
-        const size_t max = cuda_graphs_max();
+    // before `graph` gets an executable: evict the least recently used others until it fits under the caps (alone, it is
+    // kept whatever its nodes); returns what is held with it
+    cuda_graphs_held cuda_graph_make_room(const ggml_cuda_graph * graph) {
+        const size_t max       = cuda_graphs_max();
+        const size_t max_nodes = cuda_graph_nodes_max();
         for (;;) {
-            size_t held = 0;
-            auto   lru  = cuda_graphs.end();
+            cuda_graphs_held held = { 1, graph->num_nodes };
+            auto             lru  = cuda_graphs.end();
             for (auto it = cuda_graphs.begin(); it != cuda_graphs.end(); ++it) {
                 if (it->second->instance == nullptr || it->second.get() == graph) {
                     continue;
                 }
-                ++held;
+                held.graphs++;
+                held.nodes += it->second->num_nodes;
                 if (lru == cuda_graphs.end() || it->second->last_used_time < lru->second->last_used_time) {
                     lru = it;
                 }
             }
-            if (max == 0 || held < max) {
-                return held + 1;
+            const bool fits = (max == 0 || held.graphs <= max) && (max_nodes == 0 || held.nodes <= max_nodes);
+            if (fits || lru == cuda_graphs.end()) {
+                return held;
             }
             cuda_graphs.erase(lru);
         }

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -2986,10 +2986,12 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx
     return res;
 }
 
-// create `graph`'s executable from its captured graph, first evicting what the cap needs (cuda_graphs_max); the device
-// memory that takes is measured, and logged whenever it or the count held is the most on this context so far
+// create `graph`'s executable from its captured graph, first evicting what the caps need (cuda_graph_make_room); the
+// device memory that takes is measured (what the driver's pool grew by, if it did), and logged whenever what is held or
+// the most one took is the most on this context so far: at INFO when that is the most memory, or the nodes held reach a
+// power of two (a -sm tensor split's context holds hundreds of small graphs), at DEBUG otherwise
 static void ggml_cuda_graph_instantiate(ggml_backend_cuda_context * cuda_ctx, ggml_cuda_graph * graph) {
-    const size_t held = cuda_ctx->cuda_graph_make_room(graph);
+    const ggml_backend_cuda_context::cuda_graphs_held held = cuda_ctx->cuda_graph_make_room(graph);
 
     size_t free_before;
     size_t free_after;
@@ -2998,13 +3000,29 @@ static void ggml_cuda_graph_instantiate(ggml_backend_cuda_context * cuda_ctx, gg
     CUDA_CHECK(cudaGraphInstantiate(&graph->instance, graph->graph, NULL, NULL, 0));
     CUDA_CHECK(cudaMemGetInfo(&free_after, &total));
     const size_t bytes = free_before > free_after ? free_before - free_after : 0;
-
-    if (held > cuda_ctx->cuda_graphs_held_max || bytes > cuda_ctx->cuda_graph_bytes_max) {
-        cuda_ctx->cuda_graphs_held_max = std::max(cuda_ctx->cuda_graphs_held_max, held);
-        cuda_ctx->cuda_graph_bytes_max = std::max(cuda_ctx->cuda_graph_bytes_max, bytes);
-        GGML_LOG_INFO("%s: %s (context %p): %zu CUDA graphs held at most (cap %zu), the largest took %.2f MiB to instantiate\n",
+    cuda_ctx->cuda_graph_bytes += bytes;
+
+    if (held.graphs > cuda_ctx->cuda_graphs_held_max || held.nodes > cuda_ctx->cuda_graph_nodes_held_max ||
+            bytes > cuda_ctx->cuda_graph_bytes_max) {
+        const auto pow2_floor = [](size_t n) {
+            size_t p = 0;
+            for (size_t q = 1; q != 0 && q <= n; q <<= 1) {
+                p = q;
+            }
+            return p;
+        };
+        const bool info = bytes > cuda_ctx->cuda_graph_bytes_max ||
+            pow2_floor(held.nodes) > pow2_floor(cuda_ctx->cuda_graph_nodes_held_max);
+        cuda_ctx->cuda_graphs_held_max      = std::max(cuda_ctx->cuda_graphs_held_max, held.graphs);
+        cuda_ctx->cuda_graph_nodes_held_max = std::max(cuda_ctx->cuda_graph_nodes_held_max, held.nodes);
+        cuda_ctx->cuda_graph_bytes_max      = std::max(cuda_ctx->cuda_graph_bytes_max, bytes);
+        ggml_log_internal(info ? GGML_LOG_LEVEL_INFO : GGML_LOG_LEVEL_DEBUG,
+                "%s: %s (context %p): %zu CUDA graphs held at most, %zu nodes (caps %zu and %zu, 0 for none); "
+                "instantiating one took %.2f MiB at most, all of them %.2f MiB\n",
                 __func__, cuda_ctx->name.c_str(), (void *) cuda_ctx, cuda_ctx->cuda_graphs_held_max,
-                ggml_backend_cuda_context::cuda_graphs_max(), cuda_ctx->cuda_graph_bytes_max / 1024.0 / 1024.0);
+                cuda_ctx->cuda_graph_nodes_held_max, ggml_backend_cuda_context::cuda_graphs_max(),
+                ggml_backend_cuda_context::cuda_graph_nodes_max(), cuda_ctx->cuda_graph_bytes_max / 1024.0 / 1024.0,
+                cuda_ctx->cuda_graph_bytes / 1024.0 / 1024.0);
     }
 }
 
@@ -5509,6 +5527,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud
             }
 
             CUDA_CHECK(cudaStreamEndCapture(cuda_ctx->stream(), &graph->graph));
+            CUDA_CHECK(cudaGraphGetNodes(graph->graph, nullptr, &graph->num_nodes));
             graph_evaluated_or_captured = true; // CUDA graph has been captured
 
             std::lock_guard<std::mutex> lock(ggml_cuda_lock);
@@ -6906,7 +6925,8 @@ ggml_backend_t ggml_backend_cuda_init(int device) {
 #ifdef USE_CUDA_GRAPH
     static std::once_flag cuda_graphs_max_logged;
     std::call_once(cuda_graphs_max_logged, [] {
-        GGML_LOG_INFO("%s: at most %zu CUDA graphs kept per context (GGML_CUDA_GRAPH_MAX; 0 = no cap)\n", __func__,
+        GGML_LOG_INFO("%s: CUDA graphs kept per context: at most %zu nodes (GGML_CUDA_GRAPH_NODES) and %zu graphs "
+                "(GGML_CUDA_GRAPH_MAX); 0 = no cap\n", __func__, ggml_backend_cuda_context::cuda_graph_nodes_max(),
                 ggml_backend_cuda_context::cuda_graphs_max());
     });
 #endif // USE_CUDA_GRAPH

tests/CMakeLists.txt

@@ -334,6 +334,7 @@ if (NOT GGML_BACKEND_DL)
     llama_build_and_test(test-cuda-graph-key.cpp)
     llama_build_and_test(test-cuda-graph-src-type.cpp)
     llama_build_and_test(test-cuda-graph-cap.cpp)
+    llama_test(test-cuda-graph-cap NAME test-cuda-graph-cap-nodes ARGS nodes)
     llama_build_and_test(test-pq2-mma-device-memory.cpp)
 llama_build_and_test(test-ptq1_0-element-map.cpp)
 llama_build_and_test(test-ptq1_0-cuda-dot.cpp)