Read each K/V byte once in gqa-8 decode attention

mlx/backend/metal/kernels/scaled_dot_product_attention.metal

@@ -28,12 +28,24 @@ using namespace metal;
       qk_dim,                                                  \
       value_dim)
 
+#define instantiate_sdpa_vector_gqa(type, qk_dim, value_dim, hpt)  \
+  instantiate_kernel(                                              \
+      "sdpa_vector_2pass_1_gqa_" #type "_" #qk_dim "_" #value_dim, \
+      sdpa_vector_2pass_1_gqa,                                     \
+      type,                                                        \
+      qk_dim,                                                      \
+      value_dim,                                                   \
+      8,                                                           \
+      hpt)
+
 #define instantiate_sdpa_vector_heads(type)      \
   instantiate_sdpa_vector(type, 64, 64)          \
   instantiate_sdpa_vector(type, 96, 96)          \
   instantiate_sdpa_vector(type, 128, 128)        \
   instantiate_sdpa_vector(type, 192, 128)        \
   instantiate_sdpa_vector(type, 256, 256)        \
+  instantiate_sdpa_vector_gqa(type, 64, 64, 8)   \
+  instantiate_sdpa_vector_gqa(type, 128, 128, 4) \
   instantiate_sdpa_vector_aggregation(type, 64)  \
   instantiate_sdpa_vector_aggregation(type, 96)  \
   instantiate_sdpa_vector_aggregation(type, 128) \

mlx/backend/metal/kernels/sdpa_vector.h

@@ -317,6 +317,149 @@ template <typename T, int D, int V = D>
   }
 }
 
+// Duplication-free variant for high gqa_factor decode: each simdgroup owns a
+// contiguous token sub-chunk and computes HPT of its group's query heads, so
+// each K/V byte is read G / HPT times instead of G times. Single-token
+// queries without mask or sinks only; the partials layout matches
+// sdpa_vector_2pass_2.
+template <typename T, int D, int V, int G, int HPT>
+[[kernel]] void sdpa_vector_2pass_1_gqa(
+    const device T* queries [[buffer(0)]],
+    const device T* keys [[buffer(1)]],
+    const device T* values [[buffer(2)]],
+    device T* out [[buffer(3)]],
+    device float* sums [[buffer(4)]],
+    device float* maxs [[buffer(5)]],
+    const constant int& N [[buffer(7)]],
+    const constant size_t& k_head_stride [[buffer(8)]],
+    const constant size_t& k_seq_stride [[buffer(9)]],
+    const constant size_t& v_head_stride [[buffer(10)]],
+    const constant size_t& v_seq_stride [[buffer(11)]],
+    const constant float& scale [[buffer(12)]],
+    uint3 tid [[threadgroup_position_in_grid]],
+    uint3 tpg [[threadgroups_per_grid]],
+    uint3 tidtg [[thread_position_in_threadgroup]],
+    uint simd_lid [[thread_index_in_simdgroup]]) {
+  constexpr int BD = 32;
+  constexpr int qk_per_thread = D / BD;
+  constexpr int v_per_thread = V / BD;
+  constexpr int NT = G / HPT;
+
+  typedef float U;
+
+  const int kv_head_idx = tid.x;
+  const int batch_idx = tid.y;
+  const int block_idx = tid.z;
+  const int blocks = tpg.z;
+  const int g = tidtg.y;
+  const int cchunk = g / NT;
+  const int h0 = (g % NT) * HPT;
+  const int num_kv_heads = tpg.x;
+  const int num_q_heads = num_kv_heads * G;
+  const int base_head = batch_idx * num_q_heads + kv_head_idx * G;
+
+  const int chunk = (N + blocks - 1) / blocks;
+  const int kstart = block_idx * chunk;
+  const int kend = min(N, kstart + chunk);
+  const int sub = (chunk + HPT - 1) / HPT;
+  const int s0 = kstart + cchunk * sub;
+  const int s1 = min(kend, s0 + sub);
+
+  const device T* kp = keys + kv_head_idx * k_head_stride + s0 * k_seq_stride +
+      simd_lid * qk_per_thread;
+  const device T* vp = values + kv_head_idx * v_head_stride +
+      s0 * v_seq_stride + simd_lid * v_per_thread;
+
+  U q[HPT][qk_per_thread];
+  for (int j = 0; j < HPT; j++) {
+    const device T* qp =
+        queries + (base_head + h0 + j) * D + simd_lid * qk_per_thread;
+    for (int i = 0; i < qk_per_thread; i++) {
+      q[j][i] = static_cast<U>(scale) * qp[i];
+    }
+  }
+
+  U max_score[HPT];
+  U sum_exp_score[HPT];
+  U o[HPT][v_per_thread];
+  for (int j = 0; j < HPT; j++) {
+    max_score[j] = Limits<U>::finite_min;
+    sum_exp_score[j] = 0;
+    for (int i = 0; i < v_per_thread; i++) {
+      o[j][i] = 0;
+    }
+  }
+
+  for (int t = s0; t < s1; t++) {
+    U kr[qk_per_thread];
+    U vr[v_per_thread];
+    for (int i = 0; i < qk_per_thread; i++) {
+      kr[i] = kp[i];
+    }
+    for (int i = 0; i < v_per_thread; i++) {
+      vr[i] = vp[i];
+    }
+    kp += k_seq_stride;
+    vp += v_seq_stride;
+    for (int j = 0; j < HPT; j++) {
+      U score = 0;
+      for (int i = 0; i < qk_per_thread; i++) {
+        score += q[j][i] * kr[i];
+      }
+      score = simd_sum(score);
+      U new_max = max(max_score[j], score);
+      U factor = fast::exp(max_score[j] - new_max);
+      U exp_score = fast::exp(score - new_max);
+      max_score[j] = new_max;
+      sum_exp_score[j] = sum_exp_score[j] * factor + exp_score;
+      for (int i = 0; i < v_per_thread; i++) {
+        o[j][i] = o[j][i] * factor + exp_score * vr[i];
+      }
+    }
+  }
+
+  threadgroup U o_sh[G * HPT * V];
+  threadgroup U se_sh[G * HPT];
+  threadgroup U mx_sh[G * HPT];
+  for (int j = 0; j < HPT; j++) {
+    int slot = (h0 + j) * HPT + cchunk;
+    U inv = sum_exp_score[j] > 0 ? 1 / sum_exp_score[j] : 0;
+    for (int i = 0; i < v_per_thread; i++) {
+      o_sh[slot * V + simd_lid * v_per_thread + i] = o[j][i] * inv;
+    }
+    if (simd_lid == 0) {
+      se_sh[slot] = sum_exp_score[j];
+      mx_sh[slot] = max_score[j];
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  U gmax = Limits<U>::finite_min;
+  for (int s = 0; s < HPT; s++) {
+    gmax = max(gmax, mx_sh[g * HPT + s]);
+  }
+  U denom = 0;
+  U acc[v_per_thread] = {0};
+  for (int s = 0; s < HPT; s++) {
+    U w = se_sh[g * HPT + s] * fast::exp(mx_sh[g * HPT + s] - gmax);
+    denom += w;
+    for (int i = 0; i < v_per_thread; i++) {
+      acc[i] += w * o_sh[(g * HPT + s) * V + simd_lid * v_per_thread + i];
+    }
+  }
+
+  const int o_offset = base_head + g;
+  device T* op =
+      out + o_offset * blocks * V + block_idx * V + simd_lid * v_per_thread;
+  for (int i = 0; i < v_per_thread; i++) {
+    op[i] = static_cast<T>(acc[i]);
+  }
+  if (simd_lid == 0) {
+    sums[o_offset * blocks + block_idx] = denom;
+    maxs[o_offset * blocks + block_idx] = gmax;
+  }
+}
+
 template <typename T, int D>
 [[kernel]] void sdpa_vector_2pass_2(
     const device T* partials [[buffer(0)]],

mlx/backend/metal/scaled_dot_product_attention.cpp

@@ -430,7 +430,13 @@ void sdpa_vector_2pass(
   // Set the kernel name
   std::string kname;
   kname.reserve(64);
-  kname += "sdpa_vector_2pass_1_";
+  kname += "sdpa_vector_2pass_1";
+  if (!mask && !sinks && q.shape(2) == 1 && q.shape(1) == 8 * k.shape(1) &&
+      q.shape(-1) == v.shape(-1) && (q.shape(-1) == 64 || q.shape(-1) == 128) &&
+      k.shape(2) >= 8192) {
+    kname += "_gqa";
+  }
+  kname += "_";
   kname += get_type_string(q.dtype());
   kname += "_";
   kname += std::to_string(q.shape(-1));

python/tests/test_fast_sdpa.py

@@ -248,6 +248,20 @@ def test_sdpa_vector(self):
             )
             self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
 
+    def test_sdpa_vector_gqa_long(self):
+        scale = 1.0
+        mx.random.seed(0)
+        for Nq, Nkv, D in [(32, 4, 128), (64, 8, 64)]:
+            for L in [8192, 8201]:
+                q = 5e-1 * mx.random.normal(shape=(1, Nq, 1, D))
+                k = 5e-1 * mx.random.normal(shape=(1, Nkv, L + 32, D))[:, :, :L]
+                v = 5e-1 * mx.random.normal(shape=(1, Nkv, L + 32, D))[:, :, :L]
+                kr = mx.repeat(k, Nq // Nkv, axis=1)
+                vr = mx.repeat(v, Nq // Nkv, axis=1)
+                ref = mlx_primitives_sdpa(q, kr, vr, scale)
+                out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
+                self.assertTrue(mx.allclose(ref, out, atol=1e-4, rtol=1e-4))
+
     def test_sdpa_fully_masked(self):
         Lkv = 8
         mask = mx.array(False)