mirror of
https://github.com/ggml-org/whisper.cpp.git
synced 2026-10-22 03:48:19 +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_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_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(fattn_kernel, launch_params,
|
||||||
(const char *) Q->data,
|
(const char *) Q->data,
|
||||||
K_data,
|
K_data,
|
||||||
V_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(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, 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, 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(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, 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(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, 8, 128, 2, 64, 128, 128, 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, 16, 256, 1, 64, 128, 128, 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, 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(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);
|
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)
|
#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_A_KQ = tile<16, 8, half2>; // row-major
|
||||||
using T_B_KQ = tile<16, 8, half2>; // column-major
|
using T_B_KQ = tile<16, 8, half2>; // column-major
|
||||||
using T_C_KQ = tile<16, 16, float>; // 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_B_VKQ = tile<16, 8, half2>; // column-major
|
||||||
using T_C_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_A_KQ = tile<16, 8, half2>; // row-major
|
||||||
using T_B_KQ = tile< 8, 8, half2>; // column-major
|
using T_B_KQ = tile< 8, 8, half2>; // column-major
|
||||||
using T_C_KQ = tile<16, 8, float>; // row-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)
|
#elif defined(AMD_WMMA_AVAILABLE)
|
||||||
#ifdef RDNA3
|
#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_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_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
|
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
|
using T_C_VKQ = tile<16, 16, float, DATA_LAYOUT_I_MAJOR>; // column-major
|
||||||
};
|
};
|
||||||
#else
|
#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_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_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
|
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
|
#endif // RDNA3
|
||||||
#elif defined(AMD_MFMA_AVAILABLE)
|
#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_A_KQ = tile<16, 8, half2>; // row-major
|
||||||
using T_B_KQ = tile<16, 8, half2>; // column-major
|
using T_B_KQ = tile<16, 8, half2>; // column-major
|
||||||
using T_C_KQ = tile<16, 16, float>; // 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
|
using T_C_VKQ = tile<16, 16, float>; // column-major
|
||||||
};
|
};
|
||||||
#else // Volta
|
#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_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_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
|
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 warp_size = ggml_cuda_get_physical_warp_size();
|
||||||
constexpr int ncols = ncols1 * ncols2;
|
constexpr int ncols = ncols1 * ncols2;
|
||||||
using T_A_KQ = typename mma_tile_sizes<DV, ncols>::T_A_KQ;
|
using T_A_KQ = typename mma_tile_sizes<DKQ, ncols>::T_A_KQ;
|
||||||
using T_B_KQ = typename mma_tile_sizes<DV, ncols>::T_B_KQ;
|
using T_B_KQ = typename mma_tile_sizes<DKQ, ncols>::T_B_KQ;
|
||||||
using T_C_KQ = typename mma_tile_sizes<DV, ncols>::T_C_KQ;
|
using T_C_KQ = typename mma_tile_sizes<DKQ, ncols>::T_C_KQ;
|
||||||
using T_A_VKQ = typename mma_tile_sizes<DV, ncols>::T_A_VKQ;
|
using T_A_VKQ = typename mma_tile_sizes<DKQ, ncols>::T_A_VKQ;
|
||||||
using T_B_VKQ = typename mma_tile_sizes<DV, ncols>::T_B_VKQ;
|
using T_B_VKQ = typename mma_tile_sizes<DKQ, ncols>::T_B_VKQ;
|
||||||
using T_C_VKQ = typename mma_tile_sizes<DV, ncols>::T_C_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_warp = T_B_KQ::I;
|
||||||
constexpr int cols_per_thread = get_cols_per_thread();
|
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 &&
|
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_may_use_sparse(K->ne[0], V->ne[0], 1, 8) &&
|
||||||
ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, 1);
|
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;
|
return BEST_FATTN_KERNEL_VEC;
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
|
|||||||
Reference in New Issue
Block a user