From 48a108d3109597d9303fa1cd1459b56ae7c4eed4 Mon Sep 17 00:00:00 2001 From: Michal Tuszynski Date: Sat, 1 Aug 2026 21:56:12 +0000 Subject: [PATCH 1/3] ggml-cuda: bypass LDS staging for K/V loads when DKQ > 128 on AMD WMMA Assisted-by: Gemini --- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 49 ++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 7f4cfd5511ff..b577914ed053 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -596,6 +596,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( flash_attn_ext_f16_load_mask (mask_h + k_VKQ_0, tile_mask, stride_mask, k_VKQ_sup, jt*ncols1, ne01); } +#if defined(AMD_WMMA_AVAILABLE) + // For large head dims, K/V bypass LDS staging below so sync mask here. + if (DKQ > 128 && (ncols2 > 1 || mask_h)) { + __syncthreads(); + } +#endif // AMD_WMMA_AVAILABLE } // For MLA K and V have the same data. @@ -606,6 +612,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if constexpr (nstages <= 1) { const int k0_diff = k0_stop - k0_start; +#if defined(AMD_WMMA_AVAILABLE) + if (DKQ <= 128) { +#endif // AMD_WMMA_AVAILABLE constexpr bool use_cp_async = nstages == 1; flash_attn_ext_f16_load_tile (K_h2 + int64_t(k_VKQ_0)*stride_K + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup); @@ -613,6 +622,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( cp_async_wait_all(); } __syncthreads(); +#if defined(AMD_WMMA_AVAILABLE) + } else { + GGML_UNUSED(k0_diff); + } +#endif // AMD_WMMA_AVAILABLE } // Calculate tile of KQ: @@ -623,6 +637,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( #pragma unroll for (int k_KQ_0 = k0_start; k_KQ_0 < k0_stop; k_KQ_0 += T_A_KQ::J) { T_A_KQ K_A; +#if defined(AMD_WMMA_AVAILABLE) + if (DKQ > 128) { + load_ldmatrix(K_A, K_h2 + int64_t(k_VKQ_0 + i_KQ_0)*stride_K + k_KQ_0, stride_K); + } else +#endif // AMD_WMMA_AVAILABLE load_ldmatrix(K_A, tile_K + i_KQ_0*stride_tile_K + (k_KQ_0 - k0_start), stride_tile_K); if constexpr (cols_per_warp == 8) { mma(KQ_C[i_KQ_00/(np*T_A_KQ::I)], K_A, Q_B[k_KQ_0/T_A_KQ::J]); @@ -649,6 +668,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int i_KQ_0 = i_KQ_00 + (threadIdx.y % np)*T_A_KQ::I; T_A_KQ K_A; +#if defined(AMD_WMMA_AVAILABLE) + if (DKQ > 128) { + load_ldmatrix(K_A, K_h2 + int64_t(k_VKQ_0 + i_KQ_0)*stride_K + k_KQ_0, stride_K); + } else +#endif // AMD_WMMA_AVAILABLE load_ldmatrix(K_A, tile_K + i_KQ_0*stride_tile_K + (k_KQ_0 - k0_start), stride_tile_K); if constexpr (cols_per_warp == 8) { @@ -668,6 +692,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } if constexpr (nstages <= 1) { +#if defined(AMD_WMMA_AVAILABLE) + if (DKQ <= 128) +#endif // AMD_WMMA_AVAILABLE __syncthreads(); // Only needed if tile_K == tile_V. } } @@ -957,6 +984,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if constexpr (nstages <= 1) { const int i0_diff = i0_stop - i0_start; +#if defined(AMD_WMMA_AVAILABLE) + if (DKQ <= 128) { +#endif // AMD_WMMA_AVAILABLE if (!V_is_K_view || i0_stop > 2*nbatch_K2) { constexpr bool use_cp_async = nstages == 1; flash_attn_ext_f16_load_tile @@ -966,8 +996,19 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } __syncthreads(); } +#if defined(AMD_WMMA_AVAILABLE) + } else { + GGML_UNUSED(i0_diff); + } +#endif // AMD_WMMA_AVAILABLE } +#if defined(AMD_WMMA_AVAILABLE) + const half2 * tile_V_i = DKQ > 128 ? + V_h2 + int64_t(k_VKQ_0)*stride_V + i0_start/2 : + (!V_is_K_view || i0_stop > 2*nbatch_K2 ? tile_V : tile_V + i0_start/2); +#else const half2 * tile_V_i = !V_is_K_view || i0_stop > 2*nbatch_K2 ? tile_V : tile_V + i0_start/2; +#endif // AMD_WMMA_AVAILABLE #if defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) #pragma unroll @@ -978,6 +1019,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const int k0 = k00 + (threadIdx.y % np)*T_A_VKQ::J; T_A_VKQ A; // Transposed in SRAM but not in registers, gets transposed on load. +#if defined(AMD_WMMA_AVAILABLE) + if (DKQ > 128) { + load_ldmatrix_trans(A, tile_V_i + 2*k0*stride_V + (i_VKQ_0 - i0_start)/2, stride_V); + } else +#endif // AMD_WMMA_AVAILABLE load_ldmatrix_trans(A, tile_V_i + 2*k0*stride_tile_V + (i_VKQ_0 - i0_start)/2, stride_tile_V); if constexpr (T_B_KQ::I == 8) { mma(VKQ_C[i_VKQ_0/T_A_VKQ::I], A, B[k00/(np*T_A_VKQ::J)]); @@ -1011,6 +1057,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( #endif // defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) if constexpr (nstages <= 1) { +#if defined(AMD_WMMA_AVAILABLE) + if (DKQ <= 128) +#endif // AMD_WMMA_AVAILABLE __syncthreads(); // Only needed if tile_K == tile_V. } } From d273e0b48f312ce53c770f654cf8ef44351943ff Mon Sep 17 00:00:00 2001 From: Michal Tuszynski Date: Sat, 1 Aug 2026 21:56:40 +0000 Subject: [PATCH 2/3] ggml-cuda: enable MMA F16 FlashAttention for DKQ <= 256 on AMD WMMA Assisted-by: Gemini --- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 6 +++--- ggml/src/ggml-cuda/fattn.cu | 9 +++++++-- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index b577914ed053..d29b746fbc4e 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -157,8 +157,8 @@ 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, 32, 128, 2, 64, 128, 128, 128, 1, true); + GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2, 64, 128, 128, 128, 1, true); GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 32, 128, 2, 32, 160, 128, 128, 1, true); GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 64, 128, 2, 32, 160, 128, 128, 1, true); @@ -1808,7 +1808,7 @@ static __global__ void flash_attn_ext_f16( #endif // __CUDA_ARCH__ == GGML_CUDA_CC_TURING #if defined(AMD_WMMA_AVAILABLE) - if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 128) { + if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 256) { NO_DEVICE_CODE; return; } diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index ab7a3b297c07..e312c9bc46bf 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -512,8 +512,13 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const } // AMD WMMA is always faster than the tile kernel if the full tile width of 16 can be utilized. - if ((amd_wmma_available(cc) && gqa_opt_applies && Q->ne[0] <= 128) && Q->ne[0] != 40 && Q->ne[0] != 72 && Q->ne[1] * gqa_ratio_eff > 8) { - return BEST_FATTN_KERNEL_MMA_F16; + if (amd_wmma_available(cc) && gqa_opt_applies && Q->ne[0] != 40 && Q->ne[0] != 72) { + if (Q->ne[0] <= 128 && Q->ne[1] * gqa_ratio_eff > 8) { + return BEST_FATTN_KERNEL_MMA_F16; + } + if (Q->ne[0] <= 256 && Q->ne[1] * gqa_ratio_eff > 64) { + return BEST_FATTN_KERNEL_MMA_F16; + } } // If there are no tensor cores available, use the generic tile kernel: From d76c0046947c8b3fe92949fffeb634bbe7cc5d40 Mon Sep 17 00:00:00 2001 From: Michal Tuszynski Date: Sun, 2 Aug 2026 19:50:07 +0000 Subject: [PATCH 3/3] ggml-cuda: fix tile_mask write-after-read race on AMD WMMA DKQ > 128 When DKQ > 128 on AMD WMMA, K/V loads bypass LDS so the end-of-loop __syncthreads() barriers were skipped entirely. However, tile_mask still transits LDS and is read during softmax. Without a barrier at the end of the VKQ section, the next iteration's load_mask can overwrite tile_mask while a slow warp is still reading it from the current iteration. Fix: keep the __syncthreads() at the end of the V loop when the mask is active (ncols2 > 1 || mask_h), even when DKQ > 128. Reported-by: williamtwomey (PR #26419 review) Tested: 5 consecutive runs of test-backend-ops -o FLASH_ATTN_EXT on gfx1201, 2920/2920 passed each time (14600/14600 total, 0 failures). --- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 8 ++++++-- tests/test-backend-ops.cpp | 3 +++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index d29b746fbc4e..bc89684c6d6f 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -1058,9 +1058,13 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( if constexpr (nstages <= 1) { #if defined(AMD_WMMA_AVAILABLE) - if (DKQ <= 128) + // When DKQ > 128, K/V bypass LDS so tile_K/tile_V barriers are unnecessary. + // However, tile_mask still lives in LDS — the next iteration's load_mask + // would overwrite it while a slow warp might still be reading it in softmax. + // Keep the barrier when mask is active to prevent this WAR hazard. + if (DKQ <= 128 || ncols2 > 1 || mask_h) #endif // AMD_WMMA_AVAILABLE - __syncthreads(); // Only needed if tile_K == tile_V. + __syncthreads(); // Needed if tile_K == tile_V, or if tile_mask is in use. } } #else diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 4098acaaf91a..29020d69c0bf 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9594,6 +9594,9 @@ static std::vector> make_test_cases_eval() { // large-KV F16 cases (Qwen3.6-27B geometry and a llama-class control): the upstream matrix // stops at kv=1024, blind to long-context FA bugs (e.g. the oneDNN SDPA ordering race on BMG). + // For DKQ>128 on AMD WMMA, K/V bypass LDS but tile_mask stays in shared memory, so the end-of-loop + // barrier must still fire when a mask is active. kv=16384 with gqa_ratio>=4 (ncols2>=4) gives 256+ + // attention loop iterations, stressing this synchronization path. for (int64_t kv : { 4096, 16384 }) { test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, kv, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));