Commit ce8caa6e6 for llama.cpp
commit ce8caa6e60a03093351d6016a818720e0d46f0fb
Author: Johannes Gäßler <johannesg@5d6.de>
Date: Sun Sep 20 22:20:12 2026 +0200
CUDA: tune FA for Gemma 4 on Ampere or newer (#29152)
diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh
index b527b1cf6..6d1ce52db 100644
--- a/ggml/src/ggml-cuda/fattn-common.cuh
+++ b/ggml/src/ggml-cuda/fattn-common.cuh
@@ -1235,8 +1235,8 @@ void launch_fattn(
GGML_ASSERT(block_dim.x % warp_size == 0);
- ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks_num, block_dim, nbytes_shared, main_stream);
- ggml_cuda_kernel_launch(fattn_kernel, launch_params,
+ ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks_num, block_dim, nbytes_shared, main_stream);
+ ggml_cuda_kernel_launch(fattn_kernel, launch_params,
(const char *) Q->data,
K_data,
V_data,
diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
index dc18a091a..df7dd8615 100644
--- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh
+++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh
@@ -68,16 +68,16 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co
GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 64, 128, 2, 32, 96, 64, 64, 2, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 8, 128, 2, 64, 128, 128, 128, 2, true);
- GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16, 64, 4, 32, 128, 128, 128, 2, true);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16, 256, 1, 64, 128, 128, 128, 2, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 128, 2, 32, 128, 128, 128, 2, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 128, 2, 32, 128, 128, 128, 2, true);
GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 32, 128, 2, 32, 128, 128, 128, 1, false);
GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 64, 256, 1, 32, 128, 128, 128, 1, false);
- GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 8, 64, 4, 32, 256, 256, 128, 1, false);
- GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 16, 64, 4, 32, 256, 256, 128, 1, false);
- GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 32, 128, 2, 32, 128, 128, 128, 1, false);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 8, 128, 2, 64, 128, 128, 128, 1, false);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 16, 256, 1, 64, 128, 128, 128, 1, false);
+ GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 32, 256, 1, 32, 128, 128, 128, 1, false);
GGML_CUDA_FATTN_MMA_CONFIG_CASE(512, 512, 64, 256, 1, 32, 128, 128, 128, 1, false);
GGML_CUDA_FATTN_MMA_CONFIG_CASE(576, 512, 8, 64, 4, 32, 288, 256, 128, 1, false);
@@ -1066,7 +1066,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
}
#if defined(TURING_MMA_AVAILABLE)
-template<int DV, int ncols> struct mma_tile_sizes {
+template<int DKQ, int ncols> struct mma_tile_sizes {
using T_A_KQ = tile<16, 8, half2>; // row-major
using T_B_KQ = tile<16, 8, half2>; // column-major
using T_C_KQ = tile<16, 16, float>; // column-major
@@ -1074,7 +1074,33 @@ template<int DV, int ncols> struct mma_tile_sizes {
using T_B_VKQ = tile<16, 8, half2>; // column-major
using T_C_VKQ = tile<16, 8, half2>; // column-major
};
-template<int DV> struct mma_tile_sizes<DV, 8> {
+// If there are only 8 columns, use thinner B tiles to avoid wasting compute:
+template<int DKQ> struct mma_tile_sizes<DKQ, 8> {
+ using T_A_KQ = tile<16, 8, half2>; // row-major
+ using T_B_KQ = tile< 8, 8, half2>; // column-major
+ using T_C_KQ = tile<16, 8, float>; // row-major
+ using T_A_VKQ = tile<16, 8, half2>; // row-major
+ using T_B_VKQ = tile< 8, 8, half2>; // column-major
+ using T_C_VKQ = tile<16, 4, half2>; // row-major
+};
+// For very large head sizes, use thinner B tiles to reduce register pressure:
+template<> struct mma_tile_sizes<256, 16> {
+ using T_A_KQ = tile<16, 8, half2>; // row-major
+ using T_B_KQ = tile< 8, 8, half2>; // column-major
+ using T_C_KQ = tile<16, 8, float>; // row-major
+ using T_A_VKQ = tile<16, 8, half2>; // row-major
+ using T_B_VKQ = tile< 8, 8, half2>; // column-major
+ using T_C_VKQ = tile<16, 4, half2>; // row-major
+};
+template<> struct mma_tile_sizes<512, 16> {
+ using T_A_KQ = tile<16, 8, half2>; // row-major
+ using T_B_KQ = tile< 8, 8, half2>; // column-major
+ using T_C_KQ = tile<16, 8, float>; // row-major
+ using T_A_VKQ = tile<16, 8, half2>; // row-major
+ using T_B_VKQ = tile< 8, 8, half2>; // column-major
+ using T_C_VKQ = tile<16, 4, half2>; // row-major
+};
+template<> struct mma_tile_sizes<512, 32> {
using T_A_KQ = tile<16, 8, half2>; // row-major
using T_B_KQ = tile< 8, 8, half2>; // column-major
using T_C_KQ = tile<16, 8, float>; // row-major
@@ -1084,7 +1110,7 @@ template<int DV> struct mma_tile_sizes<DV, 8> {
};
#elif defined(AMD_WMMA_AVAILABLE)
#ifdef RDNA3
-template<int DV, int ncols> struct mma_tile_sizes {
+template<int DKQ, int ncols> struct mma_tile_sizes {
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
@@ -1109,7 +1135,7 @@ template<int ncols> struct mma_tile_sizes<112, ncols> {
using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>; // column-major
};
#else
-template<int DV, int ncols> struct mma_tile_sizes {
+template<int DKQ, int ncols> struct mma_tile_sizes {
using T_A_KQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR>; // row-major
using T_B_KQ = tile<16, 8, half2, DATA_LAYOUT_I_MAJOR>; // column-major
using T_C_KQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>; // column-major
@@ -1135,7 +1161,7 @@ template<int ncols> struct mma_tile_sizes<112, ncols> {
};
#endif // RDNA3
#elif defined(AMD_MFMA_AVAILABLE)
-template<int DV, int ncols> struct mma_tile_sizes {
+template<int DKQ, int ncols> struct mma_tile_sizes {
using T_A_KQ = tile<16, 8, half2>; // row-major
using T_B_KQ = tile<16, 8, half2>; // column-major
using T_C_KQ = tile<16, 16, float>; // column-major
@@ -1144,7 +1170,7 @@ template<int DV, int ncols> struct mma_tile_sizes {
using T_C_VKQ = tile<16, 16, float>; // column-major
};
#else // Volta
-template<int DV, int ncols> struct mma_tile_sizes {
+template<int DKQ, int ncols> struct mma_tile_sizes {
using T_A_KQ = tile< 8, 4, half2, DATA_LAYOUT_I_MAJOR_MIRRORED>; // row-major
using T_B_KQ = tile<32, 4, half2, DATA_LAYOUT_I_MAJOR>; // column-major
using T_C_KQ = tile<32, 8, float, DATA_LAYOUT_I_MAJOR>; // column-major
@@ -1185,12 +1211,12 @@ 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;
+ using T_A_KQ = typename mma_tile_sizes<DKQ, ncols>::T_A_KQ;
+ using T_B_KQ = typename mma_tile_sizes<DKQ, ncols>::T_B_KQ;
+ using T_C_KQ = typename mma_tile_sizes<DKQ, ncols>::T_C_KQ;
+ using T_A_VKQ = typename mma_tile_sizes<DKQ, ncols>::T_A_VKQ;
+ using T_B_VKQ = typename mma_tile_sizes<DKQ, ncols>::T_B_VKQ;
+ using T_C_VKQ = typename mma_tile_sizes<DKQ, ncols>::T_C_VKQ;
constexpr int cols_per_warp = T_B_KQ::I;
constexpr int cols_per_thread = get_cols_per_thread();
diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index f78e65277..7098c8b4c 100644
--- a/ggml/src/ggml-cuda/fattn.cu
+++ b/ggml/src/ggml-cuda/fattn.cu
@@ -630,7 +630,8 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
const bool sparse_decode = gqa_opt_applies && gqa_ratio > 4 &&
ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(K->ne[0], V->ne[0], 1, 8) &&
ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, 1);
- if (cc >= GGML_CUDA_CC_ADA_LOVELACE && Q->ne[1] == 1 && Q->ne[3] == 1 && !(gqa_ratio > 4 && K->ne[1] >= 8192) && !sparse_decode) {
+ if (!sparse_decode && cc >= GGML_CUDA_CC_ADA_LOVELACE && Q->ne[1] == 1 && Q->ne[3] == 1 &&
+ !(gqa_ratio > 4 && (Q->ne[0] >= 256 || K->ne[1] >= 8192))) {
return BEST_FATTN_KERNEL_VEC;
}
} else {