fix

Author: orca-zhang

src/llama.cpp

@@ -604,17 +604,13 @@ static struct ggml_tensor * llm_build_kqv(
         ggml_flash_attn_ext_set_prec(cur, GGML_PREC_F32);
 
         if (n_embd_head_v < n_embd_head_k) {
-            LLAMA_LOG_INFO("cur shape: [%ld, %ld, %ld]\n", cur->ne[0], cur->ne[1], cur->ne[2]);
-            cur = ggml_reshape_3d(ctx, cur, n_head, n_tokens, n_embd_head_v_out);
-            LLAMA_LOG_INFO("cur shape: [%ld, %ld, %ld]\n", cur->ne[0], cur->ne[1], cur->ne[2]);
-            cur = ggml_cont(ctx, ggml_view_3d(ctx, cur, n_head, n_tokens, n_embd_head_v,
-                               ggml_element_size(cur) * n_head,
-                               ggml_element_size(cur) * n_head * n_tokens,
+            cur = ggml_reshape_2d(ctx, ggml_cont(ctx, cur), n_embd_head_v_out*n_head, n_tokens);
+            cur = ggml_cont(ctx, ggml_view_2d(ctx, ggml_cont(ctx, cur), n_embd_head_v*n_head, n_tokens,
+                               ggml_element_size(cur) * n_embd_head_v_out,
                                0));
-            LLAMA_LOG_INFO("cur shape: [%ld, %ld, %ld]\n", cur->ne[0], cur->ne[1], cur->ne[2]);
+        } else {
+            cur = ggml_reshape_2d(ctx, cur, n_embd_head_v*n_head, n_tokens);
         }
-
-        cur = ggml_reshape_2d(ctx, cur, n_embd_head_v*n_head, n_tokens);
     } else {
         struct ggml_tensor * kq = ggml_mul_mat(ctx, k, q);
         cb(kq, "kq", il);