Update paged-attention-metal/attention/paged_attention.metal

paged-attention-metal/attention/paged_attention.metal

@@ -683,21 +683,26 @@ inline Bfloat8_ fp8_convert<Bfloat8_, Uchar8_>(const thread Uchar8_ &in,
 // TODO(EricLBuehler): optimize with vectorization
 template <int THREAD_GROUP_SIZE, typename Vec, int N>
 inline float qk_dot_(const threadgroup Vec (&q)[N], const thread Vec (&k)[N]) {
-  // Compute the parallel products for Q*K^T (treat vector lanes separately).
   using A_vec = typename FloatVec<Vec>::Type;
   A_vec qk_vec = mul<A_vec, Vec, Vec>(q[0], k[0]);
+
 #pragma unroll
   for (int ii = 1; ii < N; ++ii) {
     qk_vec = fma(q[ii], k[ii], qk_vec);
   }
 
-  // Finalize the reduction across lanes.
-  float qk = sum(qk_vec);
+  // Sum across elements within each thread
+  float qk_local = sum(qk_vec);
+
+  // Reduction within SIMD group only
 #pragma unroll
-  for (int mask = THREAD_GROUP_SIZE / 2; mask >= 1; mask /= 2) {
-    qk += simd_shuffle_xor(qk, mask);
+  for (int mask = simdgroup_size / 2; mask >= 1; mask /= 2) {
+    qk_local += simd_shuffle_xor(qk_local, mask);
   }
-  return qk;
+
+  // Each lane now has the sum across the SIMD group
+  // Return only from lane 0 or use simd_shuffle if needed
+  return simd_shuffle(qk_local, 0);
 }
 
 template <typename T, int THREAD_GROUP_SIZE> struct Qk_dot {