ggml-cuda: fix divergent barrier in f16 flash attention (#27870)

* ggml-cuda: fix divergent barrier in f16 flash attention

* ggml-cuda: avoid duplicate metadata pointer setup

ggml/src/ggml-cuda/fattn-mma-f16.cuh

@@ -1545,22 +1545,25 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
         }
     }
 
-    if (np > 1 && threadIdx.y % np == 0) {
+    if (np > 1) {
-        // Combine the meta data for parallel warps via shared memory.
-        // Warps with threadIdx.y % np != 0 must NOT return early.
-        // All threads must return simultaneously to avoid race conditions with work on the next tile.
-
         constexpr int nmeta = np*cols_per_warp >= warp_size ? np*cols_per_warp/warp_size : 1;
 
+        float KQ_cmn;
+        float KQ_cms[nmeta];
+        float KQ_crs;
+
         const int jc_meta = threadIdx.y*cols_per_warp + (np*cols_per_warp < warp_size ? threadIdx.x % (np*cols_per_warp) : threadIdx.x);
         float2 * const meta_ptr = ((float2 *) tile_Q) + jc_meta*(tile_stride/2) + nbatch_combine/2;
+
+        if (threadIdx.y % np == 0) {
+            // Combine the meta data for parallel warps via shared memory.
             float2 meta[nmeta];
 #pragma unroll
             for (int imeta = 0; imeta < nmeta; ++imeta) {
                 meta[imeta] = meta_ptr[imeta * warp_size * tile_stride/2];
             }
 
-        float KQ_cmn = meta[0].x; // KQ combine max new, max between all parallel warps.
+            KQ_cmn = meta[0].x; // KQ combine max new, max between all parallel warps.
 #pragma unroll
             for (int imeta = 1; imeta < nmeta; ++imeta) {
                 KQ_cmn = fmaxf(KQ_cmn, meta[imeta].x);
@@ -1572,13 +1575,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
                 }
             }
 
-        float KQ_cms[nmeta]; // KQ combine max scale per warp.
 #pragma unroll
             for (int imeta = 0; imeta < nmeta; ++imeta) {
                 KQ_cms[imeta] = expf(meta[imeta].x - KQ_cmn);
             }
 
-        float KQ_crs = KQ_cms[0]*meta[0].y; // KQ combine rowsum, scaled sum of all parallel warps.
+            KQ_crs = KQ_cms[0]*meta[0].y; // KQ combine rowsum, scaled sum of all parallel warps.
 #pragma unroll
             for (int imeta = 1; imeta < nmeta; ++imeta) {
                 KQ_crs += KQ_cms[imeta]*meta[imeta].y;
@@ -1589,9 +1591,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
                     KQ_crs += __shfl_xor_sync(0xFFFFFFFF, KQ_crs, offset, warp_size);
                 }
             }
+        }
 
         __syncthreads();
 
+        if (threadIdx.y % np == 0) {
             // Write back combined meta data:
 #pragma unroll
             for (int imeta = 0; imeta < nmeta; ++imeta) {
@@ -1611,11 +1615,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
                 float2 * dstk_fixup_meta = dstk_fixup + (gridDim.x + blockIdx.x)*ncols;
                 dstk_fixup_meta[(threadIdx.y/np)*cols_per_warp + threadIdx.x] = make_float2(KQ_cmn, KQ_crs);
             }
-    } else if (np > 1) {
+        }
-        // Warps with threadIdx.y % np == 0 execute a __syncthreads() in the if branch.
-        // Therefore, all other warps also need to execute a __syncthreads().
-        // Otherwise the points at which warps synchronize with each other would become misaligned.
-        __syncthreads();
     }
 
 #pragma unroll