build_rwkv6: Simplify graph

Signed-off-by: Molly Sophia <[email protected]>
ggerganov · Aug 13, 2024 · bbe7b22 · bbe7b22
1 parent 7f9de4a
commit bbe7b22
Showing 1 changed file with 12 additions and 34 deletions.
diff --git a/src/llama.cpp b/src/llama.cpp
@@ -8592,40 +8592,18 @@ static struct ggml_tensor * llm_build_time_mix_rwkv6(
         xxx
     );
 
-    struct ggml_tensor *mw = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embed * n_tokens);
-    mw = ggml_reshape_2d(
-        ctx,
-        ggml_set_1d(ctx, mw, ggml_view_1d(ctx, xxx, n_embed * n_tokens, 0), 0),
-        n_embed, n_tokens
-    );
-
-    struct ggml_tensor *mk = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embed * n_tokens);
-    mk = ggml_reshape_2d(
-        ctx,
-        ggml_set_1d(ctx, mk, ggml_view_1d(ctx, xxx, n_embed * n_tokens, n_embed * n_tokens * sizeof(float)), 0),
-        n_embed, n_tokens
-    );
-
-    struct ggml_tensor *mv = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embed * n_tokens);
-    mv = ggml_reshape_2d(
-        ctx,
-        ggml_set_1d(ctx, mv, ggml_view_1d(ctx, xxx, n_embed * n_tokens, n_embed * n_tokens * 2 * sizeof(float)), 0),
-        n_embed, n_tokens
-    );
-
-    struct ggml_tensor *mr = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embed * n_tokens);
-    mr = ggml_reshape_2d(
-        ctx,
-        ggml_set_1d(ctx, mr, ggml_view_1d(ctx, xxx, n_embed * n_tokens, n_embed * n_tokens * 3 * sizeof(float)), 0),
-        n_embed, n_tokens
-    );
-
-    struct ggml_tensor *mg = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embed * n_tokens);
-    mg = ggml_reshape_2d(
-        ctx,
-        ggml_set_1d(ctx, mg, ggml_view_1d(ctx, xxx, n_embed * n_tokens, n_embed * n_tokens * 4 * sizeof(float)), 0),
-        n_embed, n_tokens
-    );
+    // struct ggml_tensor *mw = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embed * n_tokens);
+    // mw = ggml_reshape_2d(
+    //     ctx,
+    //     ggml_set_1d(ctx, mw, ggml_view_1d(ctx, xxx, n_embed * n_tokens, 0), 0),
+    //     n_embed, n_tokens
+    // );
+
+    struct ggml_tensor *mw = ggml_view_2d(ctx, xxx, n_embed, n_tokens, xxx->nb[1], 0);
+    struct ggml_tensor *mk = ggml_view_2d(ctx, xxx, n_embed, n_tokens, xxx->nb[1], n_embed * n_tokens * sizeof(float));
+    struct ggml_tensor *mv = ggml_view_2d(ctx, xxx, n_embed, n_tokens, xxx->nb[1], n_embed * n_tokens * 2 * sizeof(float));
+    struct ggml_tensor *mr = ggml_view_2d(ctx, xxx, n_embed, n_tokens, xxx->nb[1], n_embed * n_tokens * 3 * sizeof(float));
+    struct ggml_tensor *mg = ggml_view_2d(ctx, xxx, n_embed, n_tokens, xxx->nb[1], n_embed * n_tokens * 4 * sizeof(float));
 
     struct ggml_tensor * xw = ggml_add(
         ctx,