CUDA: tune FA for Gemma 4 on Ampere or newer (llama/29152)

This commit is contained in:
Johannes Gäßler
2026-09-23 20:46:47 +03:00
committed by Georgi Gerganov
parent 984e400cc0
commit 4ef9fd70b4
3 changed files with 46 additions and 19 deletions
+2 -2
View File
@@ -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,
+42 -16
View File
@@ -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();
+2 -1
View File
@@ -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 {