mirror of
https://github.com/ggml-org/whisper.cpp.git
synced 2026-09-30 11:36:38 +02:00
CUDA: tune FA for Gemma 4 on Ampere or newer (llama/29152)
This commit is contained in:
committed by
Georgi Gerganov
parent
984e400cc0
commit
4ef9fd70b4
@@ -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,
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user