Add back MQA 2D x 2D optimisation

Author: jukofyork

src/llama-graph.cpp

@@ -1203,9 +1203,10 @@ ggml_tensor * llm_graph_context::build_attn_mha(
     // note: for MLA with the absorption optimization, the final embedding size will be changed via v_mla
     const auto n_embd_head_v = v_mla == nullptr ? v_trans ? v->ne[1] : v->ne[0] : v_mla->ne[1];
 
-    const auto n_tokens = q->ne[1];
-    const auto n_head   = q->ne[2];
-    const auto n_kv     = k->ne[1];
+    const auto n_tokens  = q->ne[1];
+    const auto n_head    = q->ne[2];
+    const auto n_kv      = k->ne[1];
+    const auto n_head_kv = k->ne[2];
 
     ggml_tensor * cur;
 
@@ -1233,12 +1234,21 @@ ggml_tensor * llm_graph_context::build_attn_mha(
 
         cur = ggml_reshape_2d(ctx0, cur, n_embd_head_v*n_head, n_tokens);
     } else {
+        // for MQA (ie: GQA with 1 group) we don't need to use a batched matrix multiply
+        if (ggml_is_contiguous(q) && n_head_kv == 1) {
+            q = ggml_reshape_2d(ctx0, q, n_embd, n_tokens*n_head);
+        }
+
         ggml_tensor * kq = ggml_mul_mat(ctx0, k, q);
 
         // note: this op tends to require high floating point range
         //       while for some models F16 is enough, for others it is not, so we default to F32 here
         ggml_mul_mat_set_prec(kq, GGML_PREC_F32);
 
+        if (ggml_is_contiguous(q) && n_head_kv == 1) {
+            kq = ggml_reshape_3d(ctx0, kq, n_kv, n_tokens, n_head);
+        }
+
         if (arch == LLM_ARCH_GROK) {
             // need to do the following:
             // multiply by attn_output_multiplyer of 0.08838834764831845