From 85beeeb11e6a6cfb66b7c2c9af6ee252aac85a89 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Sun, 5 Jul 2026 09:05:01 -0700 Subject: [PATCH] ggml-cuda: use mma-f16 flash-attn for D=256 on RDNA3.5 Route D=256 GQA attention to the mma-f16 kernel and allow its AMD_WMMA path up to DKQ=256. On RDNA the kernel still runs on WMMA instructions, and it packs ncols2 GQA heads into a full 16-wide tile that the wmma kernel leaves mostly empty. Assisted-by: Claude Opus 4 (1M context) --- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 5 +++++ ggml/src/ggml-cuda/fattn.cu | 7 +++++++ 2 files changed, 12 insertions(+) diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 83478a02cb61..e12360c07a66 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -1759,7 +1759,12 @@ static __global__ void flash_attn_ext_f16( #endif // __CUDA_ARCH__ == GGML_CUDA_CC_TURING #if defined(AMD_WMMA_AVAILABLE) + // DKQ=256 is only tuned/validated on RDNA3.5; other AMD WMMA archs keep the DKQ<=128 limit. +#if defined(RDNA3_5) + if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 256) { +#else if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 128) { +#endif // defined(RDNA3_5) NO_DEVICE_CODE; return; } diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index d6c501b1d7ea..ce4a5ccded1e 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -499,6 +499,13 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const if (can_use_vector_kernel && Q->ne[1] <= 2) { return BEST_FATTN_KERNEL_VEC; } + // The mma-f16 kernel also lowers to RDNA WMMA instructions here, and at D=256 it packs + // ncols2 GQA query heads into a full 16-wide tile where the wmma kernel leaves it mostly + // empty. Needs GQA (ncols2 >= 2) and a filled tile (ncols1*ncols2 >= 16). Restricted to + // RDNA3.5, the only AMD arch this was tuned/validated on. + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && Q->ne[0] == 256 && gqa_opt_applies && Q->ne[1]*gqa_ratio_eff >= 16) { + return BEST_FATTN_KERNEL_MMA_F16; + } return BEST_FATTN_KERNEL_WMMA_F16; }