Skip to content

ggml-cuda: enable MMA FlashAttention for head dim 256 on AMD RDNA #26419

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
srgtuszy wants to merge 3 commits into
base: master
Choose a base branch
from
+67 −6
Open

ggml-cuda: enable MMA FlashAttention for head dim 256 on AMD RDNA #26419

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
48a108d
ggml-cuda: bypass LDS staging for K/V loads when DKQ > 128 on AMD WMMA
srgtuszy Aug 1, 2026
d273e0b
ggml-cuda: enable MMA F16 FlashAttention for DKQ <= 256 on AMD WMMA
srgtuszy Aug 1, 2026
d76c004
ggml-cuda: fix tile_mask write-after-read race on AMD WMMA DKQ > 128
srgtuszy Aug 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 57 additions & 4 deletions ggml/src/ggml-cuda/fattn-mma-f16.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -596,6 +596,12 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
flash_attn_ext_f16_load_mask<ncols1, nwarps, nbatch_fa, use_cp_async, oob_check>
(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.
Expand All @@ -606,13 +612,21 @@ 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<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check>
(K_h2 + int64_t(k_VKQ_0)*stride_K + k0_start, tile_K, k0_diff, stride_K, k_VKQ_sup);
if (use_cp_async) {
cp_async_wait_all();
}
__syncthreads();
#if defined(AMD_WMMA_AVAILABLE)
} else {
GGML_UNUSED(k0_diff);
}
#endif // AMD_WMMA_AVAILABLE
}

// Calculate tile of KQ:
Expand All @@ -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]);
Expand All @@ -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) {
Expand All @@ -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.
}
}
Expand Down Expand Up @@ -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<stride_tile_V, nwarps, nbatch_fa, use_cp_async, oob_check>
Expand All @@ -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
Expand All @@ -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)]);
Expand Down Expand Up @@ -1011,7 +1057,14 @@ 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) {
__syncthreads(); // Only needed if tile_K == tile_V.
#if defined(AMD_WMMA_AVAILABLE)
// 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(); // Needed if tile_K == tile_V, or if tile_mask is in use.
}
}
#else
Expand Down Expand Up @@ -1759,7 +1812,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;
}
Expand Down
9 changes: 7 additions & 2 deletions ggml/src/ggml-cuda/fattn.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
3 changes: 3 additions & 0 deletions tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9594,6 +9594,9 @@ static std::vector<std::unique_ptr<test_case>> 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));
Expand Down