llama: optimize RWKV7 inference by fusing some graph operators

ggml/include/ggml-rpc.h

@@ -8,10 +8,10 @@ extern "C" {
 
 #define RPC_PROTO_MAJOR_VERSION    4
 #define RPC_PROTO_MINOR_VERSION    0
-#define RPC_PROTO_PATCH_VERSION    1
+#define RPC_PROTO_PATCH_VERSION    2
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16

ggml/include/ggml.h

@@ -563,6 +563,8 @@ extern "C" {
         GGML_OP_WIN_UNPART,
         GGML_OP_GET_REL_POS,
         GGML_OP_ADD_REL_POS,
+        GGML_OP_RWKV_LERP,
+        GGML_OP_RWKV_RK,
         GGML_OP_RWKV_WKV6,
         GGML_OP_GATED_LINEAR_ATTN,
         GGML_OP_RWKV_WKV7,
@@ -2501,6 +2503,20 @@ extern "C" {
             struct ggml_tensor  * pw,
             struct ggml_tensor  * ph);
 
+    GGML_API struct ggml_tensor * ggml_rwkv_lerp(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * x_prev,
+            struct ggml_tensor  * cur,
+            struct ggml_tensor  * weight);
+
+    GGML_API struct ggml_tensor * ggml_rwkv_rk(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * cur,
+            struct ggml_tensor  * k,
+            struct ggml_tensor  * r,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * r_k);
+
     GGML_API struct ggml_tensor * ggml_rwkv_wkv6(
             struct ggml_context * ctx,
             struct ggml_tensor  * k,
@@ -2525,8 +2541,8 @@ extern "C" {
             struct ggml_tensor  * w,
             struct ggml_tensor  * k,
             struct ggml_tensor  * v,
+            struct ggml_tensor  * kk,
             struct ggml_tensor  * a,
-            struct ggml_tensor  * b,
             struct ggml_tensor  * state);
 
     /* Solves a specific equation of the form Ax=B, where A is a triangular matrix

ggml/src/ggml-backend-meta.cpp

@@ -975,6 +975,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             case GGML_OP_WIN_UNPART:
             case GGML_OP_GET_REL_POS:
             case GGML_OP_ADD_REL_POS:
+            case GGML_OP_RWKV_LERP:
+            case GGML_OP_RWKV_RK:
             case GGML_OP_RWKV_WKV6:
             case GGML_OP_GATED_LINEAR_ATTN:
             case GGML_OP_RWKV_WKV7:

ggml/src/ggml-cpu/ggml-cpu.c

@@ -2031,6 +2031,14 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_add_rel_pos(params, tensor);
             } break;
+        case GGML_OP_RWKV_LERP:
+            {
+                ggml_compute_forward_rwkv_lerp(params, tensor);
+            } break;
+        case GGML_OP_RWKV_RK:
+            {
+                ggml_compute_forward_rwkv_rk(params, tensor);
+            } break;
         case GGML_OP_RWKV_WKV6:
             {
                 ggml_compute_forward_rwkv_wkv6(params, tensor);
@@ -2374,6 +2382,11 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
             {
                 n_tasks = n_threads;
             } break;
+        case GGML_OP_RWKV_LERP:
+        case GGML_OP_RWKV_RK:
+            {
+                n_tasks = n_threads;
+            } break;
         case GGML_OP_RWKV_WKV6:
         case GGML_OP_GATED_LINEAR_ATTN:
         case GGML_OP_RWKV_WKV7:
@@ -2993,6 +3006,64 @@ static int ggml_cpu_try_fuse_ops(
 
     struct ggml_tensor * node = cgraph->nodes[node_n];
 
+    if (node->op == GGML_OP_ADD) {
+        // ADD + MUL fusion
+        const enum ggml_op fuse_ops[] = { GGML_OP_ADD, GGML_OP_MUL };
+        if (ggml_can_fuse(cgraph, node_n, fuse_ops, 2)) {
+            struct ggml_tensor * mul_node = cgraph->nodes[node_n + 1];
+            const struct ggml_tensor * scale = (mul_node->src[0] == node)
+                ? mul_node->src[1] : mul_node->src[0];
+
+            if (node->src[0]->type == GGML_TYPE_F32 &&
+                node->src[1]->type == GGML_TYPE_F32 &&
+                node->type         == GGML_TYPE_F32 &&
+                scale->type        == GGML_TYPE_F32 &&
+                mul_node->type     == GGML_TYPE_F32 &&
+                ggml_are_same_shape(node->src[0], mul_node) &&
+                ggml_are_same_shape(node->src[1], mul_node) &&
+                ggml_are_same_shape(scale, mul_node)        &&
+                ggml_is_contiguous(node->src[0])        &&
+                ggml_is_contiguous(node->src[1])        &&
+                ggml_is_contiguous(scale)               &&
+                ggml_is_contiguous(mul_node)) {
+
+                ggml_compute_forward_add_mul_fused(params, node, mul_node);
+                return 1;
+            }
+        }
+    }
+
+    if (node->op == GGML_OP_NORM) {
+        // NORM + MUL + ADD fusion
+        const enum ggml_op fuse_ops[] = { GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD };
+        if (ggml_can_fuse(cgraph, node_n, fuse_ops, 3)) {
+            struct ggml_tensor * mul_node = cgraph->nodes[node_n + 1];
+            struct ggml_tensor * add_node = cgraph->nodes[node_n + 2];
+            const struct ggml_tensor * mul_w = (mul_node->src[0] == node)
+                ? mul_node->src[1] : mul_node->src[0];
+            const struct ggml_tensor * add_b = (add_node->src[0] == mul_node)
+                ? add_node->src[1] : add_node->src[0];
+
+            if (node->src[0]->type  == GGML_TYPE_F32 &&
+                node->type          == GGML_TYPE_F32 &&
+                mul_node->type      == GGML_TYPE_F32 &&
+                add_node->type      == GGML_TYPE_F32 &&
+                mul_w->type         == GGML_TYPE_F32 &&
+                add_b->type         == GGML_TYPE_F32 &&
+                mul_w->ne[0]        == node->ne[0]   &&
+                add_b->ne[0]        == node->ne[0]   &&
+                ggml_are_same_shape(node, add_node)  &&
+                ggml_is_contiguous_rows(node->src[0]) &&
+                ggml_is_contiguous_rows(mul_w)       &&
+                ggml_is_contiguous_rows(add_b)       &&
+                ggml_is_contiguous_rows(add_node)) {
+
+                ggml_compute_forward_norm_mul_add_fused(params, node, mul_node, add_node);
+                return 2;
+            }
+        }
+    }
+
     if (node->op == GGML_OP_RMS_NORM) {
         // RMS_NORM + MUL fusion
         const enum ggml_op fuse_ops[] = { GGML_OP_RMS_NORM, GGML_OP_MUL };