ggml-cuda.cu

@@ -2405,42 +2405,28 @@ GGML_CALL static void ggml_backend_cuda_synchronize(ggml_backend_t backend) {
     GGML_UNUSED(backend);
 }
 
-#if (CUDART_VERSION >= 12000)
-#define USE_CUDA_GRAPH
-#endif
-
-#ifdef USE_CUDA_GRAPH
-#define MAX_NODES_IN_CUDA_GRAPH 10000
 struct ggml_cudaGraph {
     int count=0;
     cudaGraph_t graph = nullptr;
     cudaGraphExec_t instance = nullptr;
     size_t numNodes = 0;
     int softmax_ne0 = 0;
-    cudaGraphNode_t nodes[MAX_NODES_IN_CUDA_GRAPH];
-    cudaKernelNodeParams params[MAX_NODES_IN_CUDA_GRAPH];
 };
-#endif
 
 GGML_CALL static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) {
     ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context;
 
     ggml_cuda_set_device(cuda_ctx->device);
 
-#ifdef USE_CUDA_GRAPH
     // Objects required for CUDA Graph
-    static ggml_cudaGraph cudaGraph;
-    bool useCudaGraph = (cudaGraph.count>=7); //avoid CUDA graphs on first few steps due to incompatible initialisations.
+#define MAX_NODES_IN_CUDA_GRAPH 10000
+    static ggml_cudaGraph cudaGraph; //TO DO move this to a suitable persistant location (and avoid use of static memory)
+    bool useCudaGraph = (cudaGraph.count>=2); //avoid CUDA graphs on first 2 steps due to incompatible initialisations.
     char** updatedKernelArg[MAX_NODES_IN_CUDA_GRAPH];
     bool cudaGraphUpdateRequired = false;
     // pointer to CUDA cpy kernel, which is required to identify
     // kernel parameters which need updated in the graph for each token
     void* ggmlCudaCpyFn = nullptr;
-
-    if(ggml_backend_cuda_get_device_count() > 1){
-        useCudaGraph = false; // disable CUDA graphs for multi-gpu for now. TO DO investigate
-    }
-
     if(useCudaGraph) {
 
         if(cudaGraph.instance == nullptr) cudaGraphUpdateRequired=true;
@@ -2452,9 +2438,6 @@ GGML_CALL static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t
             // Identify if the graph needs updated for this token due to the number of elements changing
             // (identified by inspecting soft max op parameters)
             if(node->op == GGML_OP_SOFT_MAX) {
-                if(node->src[1]->ne[1] > 1){
-                    useCudaGraph = false; // disable CUDA graphs for batch size > 1 for now. TO DO investigate
-                }
                 if(node->src[0]->ne[0] != cudaGraph.softmax_ne0) {
                     cudaGraphUpdateRequired = true;
                     cudaGraph.softmax_ne0 = node->src[0]->ne[0];
@@ -2475,11 +2458,6 @@ GGML_CALL static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t
         CUDA_CHECK(cudaStreamBeginCapture(cuda_ctx->stream(), cudaStreamCaptureModeGlobal));
     }
 
-#else
-    bool useCudaGraph = false;
-    bool cudaGraphUpdateRequired = false;
-#endif
-    
     // Only perfom the graph exection if CUDA graphs are not enebled, or we are capturing the graph.
     // With use of CUDA graphs, the execution will be performed by the graph launch.
     if(!useCudaGraph || cudaGraphUpdateRequired) {
@@ -2508,7 +2486,6 @@ GGML_CALL static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t
     }
     }
 
-    #ifdef USE_CUDA_GRAPH
     if(useCudaGraph && (cudaGraphUpdateRequired)) { // End CUDA graph capture
         CUDA_CHECK(cudaStreamEndCapture(cuda_ctx->stream(), &cudaGraph.graph));
     }
@@ -2521,56 +2498,53 @@ GGML_CALL static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t
 
         // Perform update to graph (if required for this token), and change copy parameter (required for every token)
 
+        cudaGraphNode_t nodes[MAX_NODES_IN_CUDA_GRAPH];
+        CUDA_KERNEL_NODE_PARAMS_v2 paramsDriver[MAX_NODES_IN_CUDA_GRAPH];
+        cudaKernelNodeParams paramsRuntime[MAX_NODES_IN_CUDA_GRAPH];
+
         if(cudaGraphUpdateRequired) {
             // Extract nodes from graph
             if(cudaGraph.numNodes == 0) {
                 CUDA_CHECK(cudaGraphGetNodes(cudaGraph.graph, nullptr, &cudaGraph.numNodes));
             }
-            CUDA_CHECK(cudaGraphGetNodes(cudaGraph.graph, cudaGraph.nodes, &cudaGraph.numNodes));
+            CUDA_CHECK(cudaGraphGetNodes(cudaGraph.graph, nodes, &cudaGraph.numNodes));
 
             // Loop over nodes, and extract kernel parameters fro each node
             for(size_t i=0; i<cudaGraph.numNodes; i++) {
-                cudaGraphNodeType nodeType;
-                CUDA_CHECK(cudaGraphNodeGetType(cudaGraph.nodes[i], &nodeType));
-                if (nodeType == cudaGraphNodeTypeKernel) {
-                    auto statRT = cudaGraphKernelNodeGetParams(cudaGraph.nodes[i], &cudaGraph.params[i]); // Get params using runtime
-                    if(statRT == cudaErrorInvalidDeviceFunction) {
-                        // Fails due to incorrect handling by CUDA runtime of CUDA BLAS node.
-                        // We don't need to update blas nodes, so clear error and move on.
-                        cudaGetLastError();
-                    }
+                // We currently get a set of params using both driver and runtime, to work around an issue (see below)
+                CU_CHECK(cuGraphKernelNodeGetParams(nodes[i], &paramsDriver[i])); // Get params using driver
+                cudaError_t statRT = cudaGraphKernelNodeGetParams(nodes[i], &paramsRuntime[i]); // Get params using runtime
+                if(statRT == 98) {
+                    // Fails due to incorrect handling by CUDA runtime of CUDA BLAS node.
+                    // We don't need to update blas nodes, so clear error and move on.
+                    cudaGetLastError();
                 }
             }
         }
 
         // Update copy kernel param (required every token)
+        // Currently uses runtime copy of params to identify copy function node,
+        // and driver copy of params to perform the update
+        // TO DO work out how to do it only using runtime copy.
         if(!cudaGraphUpdateRequired) { // on update steps, the live parameters will already be captured
             int k=0;
             for(size_t i=0; i<cudaGraph.numNodes; i++) {
-                if(cudaGraph.params[i].func == ggmlCudaCpyFn) {
+                if(paramsRuntime[i].func == ggmlCudaCpyFn) {
                     char** updatedKernelArgPointer = updatedKernelArg[k++];
-                    cudaGraph.params[i].kernelParams[1] = updatedKernelArgPointer;
-                    CUDA_CHECK(cudaGraphKernelNodeSetParams(cudaGraph.nodes[i], &cudaGraph.params[i]));
+                    paramsDriver[i].kernelParams[1] = updatedKernelArgPointer;
+                    CU_CHECK(cuGraphKernelNodeSetParams(nodes[i], &paramsDriver[i]));
                 }
             }
         }
 
         // Update graph executable
         cudaGraphExecUpdateResultInfo resultInfo;
-        auto stat = cudaGraphExecUpdate(cudaGraph.instance, cudaGraph.graph, &resultInfo);
-        if(stat == cudaErrorGraphExecUpdateFailure)
-        {
-            // The pre-existing graph exec cannot be updated due to violated constraints
-            // so instead clar error and re-instantiate
-            cudaGetLastError();
-            CUDA_CHECK(cudaGraphInstantiate(&cudaGraph.instance, cudaGraph.graph, NULL, NULL, 0));          
-        }
+        CUDA_CHECK(cudaGraphExecUpdate(cudaGraph.instance, cudaGraph.graph, &resultInfo));
 
         // Launch graph
         CUDA_CHECK(cudaGraphLaunch(cudaGraph.instance, cuda_ctx->stream()));
     }
     cudaGraph.count++;
-#endif
     return GGML_STATUS_SUCCESS;
 }