Make KVarN usable on ROCm RDNA3 (toward #122: D256 race, portable prefill, f32 WMMA prototype)

docs/beellama-args.md

@@ -51,6 +51,14 @@ peak transient scratch for concurrent long prompts but adds partial-softmax
 merges and changes floating-point reduction order. It does not alter context or
 persistent KV-cache capacity.
 
+On HIP/ROCm, KVarN prompt prefill defaults to the F32-accumulator WMMA route
+on arches whose tiles accumulate in fp32 (RDNA3/gfx11); RDNA4 stays on the
+portable route until its fp32 tiles qualify. `GGML_KVARN_AMD_PROMPT_PORTABLE`
+opts a prompt back into portable-native direct-record attention: any nonzero
+value (conventionally `1`) selects portable, while unset, `0`, or
+non-numeric values keep the WMMA default. The check runs before the generic
+probe, so opting in does not pay for a discarded WMMA pass.
+
 ## KV cache precision tail for quantized caches
 
 The KV cache precision tail (KVCPT) makes the newest attention-visible entries exact in F16 or BF16 for

docs/beellama-features.md

@@ -190,7 +190,8 @@ correctness, memory behavior, or performance on that GPU.
 
 | HIP architecture | Physical wave | Native KVarN route |
 |---|---:|---|
-| RDNA3, RDNA3.5, RDNA4 | 32 | WMMA generic/prefill and occupancy-selected split decode |
+| RDNA3, RDNA3.5 | 32 | WMMA generic/prefill (D256 on fp32-accumulator tiles, qualified on gfx1100) and occupancy-selected split decode |
+| RDNA4 | 32 | WMMA generic/prefill up to D128; D256+ stays on portable direct-record attention until its fp32 tiles qualify |
 | CDNA1-CDNA4 | 64 | MFMA generic/prefill and physical-wave split decode |
 | Older GCN, RDNA1, RDNA2 | device default | Portable direct-record attention |
 

ggml/src/ggml-cuda/fattn-kvarn-dispatch.cu

@@ -880,6 +880,12 @@ static bool ggml_cuda_flash_attn_ext_kvarn_decode(
 }
 
 static ggml_cuda_fattn_kvarn_amd_mma_arch ggml_cuda_fattn_kvarn_amd_arch(int cc) {
+    if (GGML_CUDA_CC_IS_RDNA4(cc)) {
+        // RDNA4 compiles the half2 WMMA tiles only: the fp32-accumulator
+        // tiles that justify the raised D256 limit are RDNA3 (gfx11) builds.
+        // Keep RDNA4 fail-closed at D128 until its fp32 tiles are qualified.
+        return GGML_CUDA_FATTN_KVARN_AMD_RDNA4_WMMA;
+    }
     if (amd_wmma_available(cc)) {
         return GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA;
     }
@@ -1225,6 +1231,31 @@ bool ggml_cuda_flash_attn_ext_kvarn(
         ggml_cuda_fattn_kvarn_portable_supported(plan, dst);
     bool generic_shape_supported = false;
     bool wide_mma = false;
+#if defined(GGML_USE_HIP)
+    // RDNA3 (gfx11) WMMA prompt tiles accumulate in fp32 for DV=128/256
+    // (mirroring the proven DV=80/112 fp32-PV tiles), so on fp32-tile arches
+    // the WMMA path is both the fast and the exact route (~1e-5 ladder RMSE,
+    // 32k KLD at portable parity). It is therefore the default for HIP KVarN
+    // prompt-prefill. RDNA4 compiles the half2 tiles only and stays
+    // fail-closed on portable (see the RDNA4 eligibility gate). Decode
+    // (nq<=16) stays on WMMA as before.
+    // Checked BEFORE the generic probe below: the probe launches the WMMA
+    // kernel to test the shape, so diverting first avoids running prompt
+    // prefill twice and discarding the WMMA pass.
+    {
+        const char * prompt_portable = getenv("GGML_KVARN_AMD_PROMPT_PORTABLE");
+        if (prompt_prefill && portable_supported &&
+                (prompt_portable != nullptr && atoi(prompt_portable) != 0)) {
+            g_kvarn_route_portable_native.fetch_add(1, std::memory_order_relaxed);
+            // QB-batching was superseded by upstream's complete optimized D64
+            // rewrite (v0.4.7); the fallback uses the standard portable kernel.
+            ggml_cuda_fattn_kvarn_debug_route(
+                ctx.device, plan, dst, entry_path, "portable-native",
+                "hip-prompt-precision-optin");
+            return ggml_cuda_flash_attn_ext_kvarn_portable(ctx, dst, plan);
+        }
+    }
+#endif
     if (capabilities.generic_mma && Q->ne[0] != 64) {
         generic_shape_supported = ggml_cuda_flash_attn_ext_mma_kvarn(ctx, dst, wide_mma);
         if (!generic_shape_supported) {

ggml/src/ggml-cuda/fattn-kvarn-portable.cuh

@@ -182,8 +182,8 @@ static __global__ void ggml_cuda_fattn_kvarn_portable_kernel(
     const float * q = (const float *) (
         q_data + query * nbq1 + query_head * nbq2 + stream * nbq3);
 
-    __shared__ float reduction[RECORD_DIM];
-    __shared__ float transform[RECORD_DIM];
+    __shared__ float reduction[D];
+    __shared__ float transform[D];
     __shared__ float maximum;
     __shared__ float denominator;
     __shared__ float old_scale_shared;
@@ -828,8 +828,7 @@ static inline bool ggml_cuda_fattn_kvarn_portable_supported(
          body_meta->ne[1] == q->ne[2] && body_meta->ne[2] == q->ne[1] &&
          body_meta->ne[3] == q->ne[3] && ggml_is_contiguous(body_meta));
     const bool domain_ok = ggml_cuda_fattn_kvarn_rotated_decode_domain(dst) ||
-        (q->ne[0] == 64 &&
-         ggml_cuda_fattn_kvarn_domain(dst) == GGML_FLASH_ATTN_EXT_KVARN_DOMAIN_ROTATED_K_ORIGINAL_V);
+        ggml_cuda_fattn_kvarn_domain(dst) == GGML_FLASH_ATTN_EXT_KVARN_DOMAIN_ROTATED_K_ORIGINAL_V;
     return domain_ok &&
         (q->ne[0] == 64 || q->ne[0] == 128 || q->ne[0] == 256 || q->ne[0] == 512) &&
         q->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 &&

ggml/src/ggml-cuda/fattn-kvarn-route-policy.h

@@ -30,6 +30,7 @@ enum ggml_cuda_fattn_kvarn_route {
 enum ggml_cuda_fattn_kvarn_amd_mma_arch {
     GGML_CUDA_FATTN_KVARN_AMD_NONE,
     GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA,
+    GGML_CUDA_FATTN_KVARN_AMD_RDNA4_WMMA,
     GGML_CUDA_FATTN_KVARN_AMD_CDNA_MFMA,
 };
 
@@ -61,14 +62,16 @@ inline ggml_cuda_fattn_kvarn_mma_eligibility ggml_cuda_fattn_kvarn_amd_mma_eligi
         return GGML_CUDA_FATTN_KVARN_MMA_INVALID_COLUMNS;
     }
     if (input.head_dim <= 0 ||
-            (input.arch == GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA && input.head_dim > 128) ||
+            (input.arch == GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA && input.head_dim > 256) ||
+            (input.arch == GGML_CUDA_FATTN_KVARN_AMD_RDNA4_WMMA && input.head_dim > 128) ||
             (input.arch == GGML_CUDA_FATTN_KVARN_AMD_CDNA_MFMA && input.head_dim > 256)) {
         return GGML_CUDA_FATTN_KVARN_MMA_HEAD_DIM_UNSUPPORTED;
     }
     if (input.ncols1 * input.ncols2 < 16) {
         return GGML_CUDA_FATTN_KVARN_MMA_TILE_TOO_SMALL;
     }
-    if (input.arch == GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA && input.ncols2 == 1) {
+    if ((input.arch == GGML_CUDA_FATTN_KVARN_AMD_RDNA_WMMA ||
+            input.arch == GGML_CUDA_FATTN_KVARN_AMD_RDNA4_WMMA) && input.ncols2 == 1) {
         return GGML_CUDA_FATTN_KVARN_MMA_RDNA_SINGLE_GQA_COLUMN;
     }
     return GGML_CUDA_FATTN_KVARN_MMA_ELIGIBLE;

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

@@ -164,7 +164,7 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256,  8,  64, 2,  32, 128, 128, 128, 1, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16,  64, 2,  32, 128, 128, 128, 1, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 128, 2,  64, 128, 128,  64, 1, true);
-    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2,  64, 128, 128,  64, 1, true);
+    GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 256, 2,  32, 128, 128,  32, 1, true);
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 128, 256, 1,  64, 128, 128,  64, 1, true);
 
     GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 32, 128, 2,  32, 160, 128, 128, 1, true);
@@ -950,12 +950,16 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
         }
 #elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
         if constexpr (std::is_same_v<decltype(T_C_VKQ::x), half2[T_C_VKQ::ne]>) {
-            const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[0], KQ_max_scale[0]);
+            // Rescale in fp32 to avoid double-rounding the scale to half first.
+            const float scale_f32 = KQ_max_scale[0];
 #pragma unroll
             for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
 #pragma unroll
                 for (int l = 0; l < T_C_VKQ::ne; ++l) {
-                    VKQ_C[i].x[l] *= KQ_max_scale_h2;
+                    float2 acc_f32 = __half22float2(VKQ_C[i].x[l]);
+                    acc_f32.x *= scale_f32;
+                    acc_f32.y *= scale_f32;
+                    VKQ_C[i].x[l] = make_half2(acc_f32.x, acc_f32.y);
                 }
             }
         } else {
@@ -1098,6 +1102,20 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
 #endif // defined(VOLTA_MMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
 }
 
+// KVarN-only tile selector. Forwards to mma_tile_sizes except RDNA3 D128/D256,
+// which use the fp32-accumulator specializations below. Lets the dense path
+// keep its qualified tiles while KVarN uses the qualified fp32 ones. The
+// member lookups are dependent and resolve after the arch regions below.
+template<int DV, int ncols> struct mma_tile_sizes;
+template<int DV, int ncols> struct mma_tile_sizes_kvarn {
+    using T_A_KQ  = typename mma_tile_sizes<DV, ncols>::T_A_KQ;
+    using T_B_KQ  = typename mma_tile_sizes<DV, ncols>::T_B_KQ;
+    using T_C_KQ  = typename mma_tile_sizes<DV, ncols>::T_C_KQ;
+    using T_A_VKQ = typename mma_tile_sizes<DV, ncols>::T_A_VKQ;
+    using T_B_VKQ = typename mma_tile_sizes<DV, ncols>::T_B_VKQ;
+    using T_C_VKQ = typename mma_tile_sizes<DV, ncols>::T_C_VKQ;
+};
+
 #if defined(TURING_MMA_AVAILABLE)
 template<int DV, int ncols> struct mma_tile_sizes {
     using T_A_KQ  = tile<16,  8, half2>; // row-major
@@ -1141,6 +1159,27 @@ template<int ncols> struct mma_tile_sizes<112, ncols> {
     using T_B_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
     using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
 };
+// KVarN-only fp32-accumulator tiles (stew675 f32-VKQ guidance): DV=128/256
+// with fp16 PV accumulator show ~3e-4/tile error compounding over 64 layers
+// on gfx1100. Selected explicitly for the KVarN path via mma_tile_sizes_kvarn
+// (forwarding primary declared above the region chain); dense keeps the
+// primary half2 tile.
+template<int ncols> struct mma_tile_sizes_kvarn<128, ncols> {
+    using T_A_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
+    using T_B_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
+    using T_C_KQ  = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
+    using T_A_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
+    using T_B_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
+    using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
+};
+template<int ncols> struct mma_tile_sizes_kvarn<256, ncols> {
+    using T_A_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
+    using T_B_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
+    using T_C_KQ  = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
+    using T_A_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
+    using T_B_VKQ = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // column-major
+    using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>;          // column-major
+};
 #else
 template<int DV, int ncols> struct mma_tile_sizes {
     using T_A_KQ  = tile<16,  8, half2, DATA_LAYOUT_I_MAJOR>;           // row-major
@@ -1220,12 +1259,15 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
 
     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
     constexpr int ncols = ncols1 * ncols2;
-    using     T_A_KQ    = typename mma_tile_sizes<DV, ncols>::T_A_KQ;
-    using     T_B_KQ    = typename mma_tile_sizes<DV, ncols>::T_B_KQ;
-    using     T_C_KQ    = typename mma_tile_sizes<DV, ncols>::T_C_KQ;
-    using     T_A_VKQ   = typename mma_tile_sizes<DV, ncols>::T_A_VKQ;
-    using     T_B_VKQ   = typename mma_tile_sizes<DV, ncols>::T_B_VKQ;
-    using     T_C_VKQ   = typename mma_tile_sizes<DV, ncols>::T_C_VKQ;
+    constexpr bool is_kvarn_kv = ggml_cuda_fattn_kvarn_template_type(type_K) || ggml_cuda_fattn_kvarn_template_type(type_V);
+    using tile_sizes_sel = typename std::conditional<is_kvarn_kv,
+        mma_tile_sizes_kvarn<DV, ncols>, mma_tile_sizes<DV, ncols>>::type;
+    using     T_A_KQ    = typename tile_sizes_sel::T_A_KQ;
+    using     T_B_KQ    = typename tile_sizes_sel::T_B_KQ;
+    using     T_C_KQ    = typename tile_sizes_sel::T_C_KQ;
+    using     T_A_VKQ   = typename tile_sizes_sel::T_A_VKQ;
+    using     T_B_VKQ   = typename tile_sizes_sel::T_B_VKQ;
+    using     T_C_VKQ   = typename tile_sizes_sel::T_C_VKQ;
 
     constexpr int  cols_per_warp   = T_B_KQ::I;
     constexpr int  cols_per_thread = get_cols_per_thread();
@@ -1235,7 +1277,6 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
     constexpr int  nbatch_V2       = ggml_cuda_fattn_mma_get_nbatch_V2     (DKQ, DV, ncols);
     constexpr int  nbatch_combine  = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols);
     constexpr bool Q_in_reg        = ggml_cuda_fattn_mma_get_Q_in_reg      (DKQ, DV, ncols);
-    constexpr bool is_kvarn_kv     = ggml_cuda_fattn_kvarn_template_type(type_K) || ggml_cuda_fattn_kvarn_template_type(type_V);
     constexpr int  nstages         = is_kvarn_kv ? 0 : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2, use_sparse);
     static_assert(!is_kvarn_kv || !use_sparse, "sparse KVarN record loads are not qualified");
 
@@ -1279,7 +1320,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
 #if defined(TURING_MMA_AVAILABLE)
     T_C_VKQ VKQ_C[cols_per_warp == 8 ? DV/T_C_VKQ::I : DV/(2*T_C_VKQ::J)];
 #elif defined(AMD_WMMA_AVAILABLE) && defined(RDNA3)
-    T_C_VKQ VKQ_C[DV % 32 != 0       ? DV/T_C_VKQ::J : DV/(2*T_C_VKQ::J)];
+    // Entry count mirrors the rescale loops: half2 accumulators fold two
+    // stacked K-halves per entry via the opsel pair (DV/32 for DV%32==0),
+    // float accumulators keep one 16-row tile per entry (DV/16 always).
+    static constexpr int VKQ_C_COUNT = std::is_same_v<decltype(T_C_VKQ::x), float[T_C_VKQ::ne]>
+        ? DV/T_C_VKQ::J
+        : (DV % 32 != 0 ? DV/T_C_VKQ::J : DV/(2*T_C_VKQ::J));
+    T_C_VKQ VKQ_C[VKQ_C_COUNT];
 #elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
     T_C_VKQ VKQ_C[                                     DV/(2*T_C_VKQ::J)];
 #else // Volta
@@ -1496,12 +1543,16 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
         }
 #elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE)
         if constexpr (std::is_same_v<decltype(T_C_VKQ::x), half2[T_C_VKQ::ne]>) {
-            const half2 KQ_max_scale_h2 = make_half2(KQ_max_scale[0], KQ_max_scale[0]);
+            // Rescale in fp32 to avoid double-rounding the scale to half first.
+            const float scale_f32 = KQ_max_scale[0];
 #pragma unroll
             for (int i = 0; i < (DV/2)/T_C_VKQ::J; ++i) {
 #pragma unroll
                 for (int l = 0; l < T_C_VKQ::ne; ++l) {
-                    VKQ_C[i].x[l] *= KQ_max_scale_h2;
+                    float2 acc_f32 = __half22float2(VKQ_C[i].x[l]);
+                    acc_f32.x *= scale_f32;
+                    acc_f32.y *= scale_f32;
+                    VKQ_C[i].x[l] = make_half2(acc_f32.x, acc_f32.y);
                 }
             }
         } else {
@@ -1562,7 +1613,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
                 float2 * dstk_fixup_meta = dstk_fixup + (gridDim.x + blockIdx.x)*ncols;
                 dstk_fixup_meta[jc_cwm] = KQ_cmr;
             }
-            if (!is_kvarn_kv && !needs_fixup && !is_fixup && dst_final_meta && threadIdx.x < T_B_KQ::I) {
+            // KVarN whole-tile blocks must publish final (max, rowsum) too: the
+            // tail merge reads body_meta for every row, and the stream-k fixup
+            // skips tiles whose K range aligns exactly to tile boundaries, so
+            // without this store those rows keep zero meta and their (correct)
+            // body values are silently discarded by the merge.
+            if (!needs_fixup && !is_fixup && dst_final_meta && threadIdx.x < T_B_KQ::I) {
                 const int j = jc_cwm / ncols2;
                 const int c = jc_cwm % ncols2;
                 if (jt*ncols1 + j < int(ne01.z) && zt_gqa*ncols2 + c < gqa_ratio) {
@@ -1608,7 +1664,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
                 float2 * dstk_fixup_meta = dstk_fixup + (gridDim.x + blockIdx.x)*ncols;
                 dstk_fixup_meta[jc_cwm] = KQ_cmr;
             }
-            if (!is_kvarn_kv && !needs_fixup && !is_fixup && dst_final_meta && thread_should_write) {
+            // KVarN whole-tile blocks must publish final (max, rowsum) too: the
+            // tail merge reads body_meta for every row, and the stream-k fixup
+            // skips tiles whose K range aligns exactly to tile boundaries, so
+            // without this store those rows keep zero meta and their (correct)
+            // body values are silently discarded by the merge.
+            if (!needs_fixup && !is_fixup && dst_final_meta && thread_should_write) {
                 const int j = jc_cwm / ncols2;
                 const int c = jc_cwm % ncols2;
                 if (jt*ncols1 + j < int(ne01.z) && zt_gqa*ncols2 + c < gqa_ratio) {
@@ -1684,7 +1745,12 @@ 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);
         }
-        if (!is_kvarn_kv && !needs_fixup && !is_fixup && dst_final_meta &&
+        // KVarN whole-tile blocks must publish final (max, rowsum) too: the
+        // tail merge reads body_meta for every row, and the stream-k fixup
+        // skips tiles whose K range aligns exactly to tile boundaries, so
+        // without this store those rows keep zero meta and their (correct)
+        // body values are silently discarded by the merge.
+        if (!needs_fixup && !is_fixup && dst_final_meta &&
                 (cols_per_warp == warp_size || threadIdx.x < cols_per_warp)) {
             const int jc = (threadIdx.y/np)*cols_per_warp + threadIdx.x;
             if (jc < ncols) {
@@ -1826,7 +1892,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
                 }
             }
         }
-        if (np > 1) {
+        // The tile_Q buffer is reused for the next k00 iteration, so all warps must sync here
+        // before its data is overwritten. With np > 1 only some warps read back, but they all write.
+        if (np > 1 || k00 + nbatch_combine < DV/2) {
             __syncthreads();
         }
     }
@@ -1915,7 +1983,7 @@ static __global__ void flash_attn_ext_f16(
 #if defined(AMD_WMMA_AVAILABLE)
     // Mirrored by ggml_cuda_fattn_kvarn_amd_mma_eligibility on the host.
     // Keep this final invariant for callers outside the KVarN dispatcher.
-    if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 128) {
+    if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 256) {
         NO_DEVICE_CODE;
         return;
     }
@@ -1998,6 +2066,11 @@ static __global__ void flash_attn_ext_f16(
                  ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop);
         }
 
+        // The next process_tile call reuses the tile_Q buffer for its Q/K tiles, so all warps must
+        // have finished reading the combined results before any of them starts the next call.
+        // (With np == 1 the end-of-k00 barrier does not fire, so this is required for correctness.)
+        __syncthreads();
+
         kbc += iter_k;
         kbc -= kbc % iter_k;