Avoid using saved CUDA graph if scale changes

Fixes ggerganov#9451
agray3 · Sep 17, 2024 · 3fe1f68 · slaren · Sep 17, 2024 · 3fe1f68
1 parent 0d2ec43
commit 3fe1f68
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu
@@ -2441,6 +2441,11 @@ GGML_CALL static void ggml_backend_cuda_synchronize(ggml_backend_t backend) {
     GGML_UNUSED(backend);
 }
 
+static int32_t ggml_get_op_params_i32(const struct ggml_tensor * tensor, uint32_t i) {
+    assert(i < GGML_MAX_OP_PARAMS / sizeof(int32_t));
+    return ((const int32_t *)(tensor->op_params))[i];
+}
+
 static void set_ggml_graph_node_properties(ggml_tensor * node, ggml_graph_node_properties * graph_node_properties) {
     graph_node_properties->node_address = node->data;
     graph_node_properties->node_op = node->op;
@@ -2451,6 +2456,9 @@ static void set_ggml_graph_node_properties(ggml_tensor * node, ggml_graph_node_p
     for (int i = 0; i < GGML_MAX_SRC; i++) {
         graph_node_properties->src_address[i] = node->src[i] ? node->src[i]->data : nullptr;
     }
+    for (int i = 0; i < GGML_MAX_OP_PARAMS; i++) {
+        graph_node_properties->op_params[i] = ggml_get_op_params_i32(node, i);
+    }
 }
 
 static bool ggml_graph_node_has_matching_properties(ggml_tensor * node, ggml_graph_node_properties * graph_node_properties) {
@@ -2482,6 +2490,14 @@ static bool ggml_graph_node_has_matching_properties(ggml_tensor * node, ggml_gra
             return false;
         }
     }
+
+    if (node->op == GGML_OP_SCALE) {
+        for (int i = 0; i < GGML_MAX_OP_PARAMS; i++) {
+            if (ggml_get_op_params_i32(node, i) != graph_node_properties->op_params[i]) {
+                return false;
+            }
+        }
+    }
     return true;
 }
 

diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh
@@ -569,6 +569,7 @@ struct ggml_graph_node_properties {
     int64_t ne[GGML_MAX_DIMS];
     size_t nb[GGML_MAX_DIMS];
     void * src_address[GGML_MAX_SRC];
+    int32_t op_params[GGML_MAX_OP_PARAMS];
 };
 
 struct ggml_cuda_graph {