public_git/whisper.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/whisper.cpp.git synced 2026-10-12 01:05:43 +02:00
Code Issues Packages Projects Releases Wiki Activity

CUDA: prefer whole-tile FlashAttention scheduling for efficient two-stage kernels (llama/29435)

Browse Source
This commit is contained in:
anujj
2026-10-06 10:38:20 +03:00
committed by Georgi Gerganov
parent ce6552ef7b
commit 8bb85d9636
2 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/fattn-common.cuh
+10 -3
View File
@@ -981,7 +981,7 @@ template <int DV, int ncols1, int ncols2>
void launch_fattn(
ggml_backend_cuda_context & ctx, ggml_tensor * dst, fattn_kernel_t fattn_kernel, const int nwarps, const size_t nbytes_shared,
const int nbatch_fa, const bool need_f16_K, const bool need_f16_V, const bool stream_k, const bool use_sparse,
const int warp_size = WARP_SIZE
const int warp_size = WARP_SIZE, const bool async_kv_preload = false
) {
constexpr int ncols = ncols1 * ncols2;
@@ -1114,7 +1114,8 @@ void launch_fattn(
// Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped.
// Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or
// multiple sequences of possibly different lengths.
if (!use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) {
const bool scan_mask = !use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1);
if (scan_mask) {
const int64_t s31 = mask->nb[1] / sizeof(half2);
const int64_t s33 = mask->nb[3] / sizeof(half2);
@@ -1142,10 +1143,16 @@ void launch_fattn(
dim3 blocks_num;
if (stream_k) {
auto should_use_stream_k = [](const int cc, const int ntiles_dst, const int max_blocks, const int DKQ) {
// Stream-K splits the work before the mask scan is applied, so skipped KV tiles make the blocks uneven.
const bool prefer_whole_tiles = GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_DGX_SPARK && async_kv_preload && scan_mask;
auto should_use_stream_k = [prefer_whole_tiles](const int cc, const int ntiles_dst, const int max_blocks, const int DKQ) {
const int tiles_nwaves = (ntiles_dst + max_blocks - 1) / max_blocks;
const int tiles_efficiency_percent = 100 * ntiles_dst / (max_blocks*tiles_nwaves);
if (prefer_whole_tiles && tiles_efficiency_percent >= 75) {
return false;
}
if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_ADA_LOVELACE) {
return true;
}
ggml/src/ggml-cuda/fattn-mma-f16.cuh
+2 -1
View File
@@ -2100,8 +2100,9 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
}
}
const bool async_kv_preload = nstages == 2 && !use_sparse;
launch_fattn<DV, ncols1, ncols2>
(ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host);
(ctx, dst, fattn_kernel, nwarps, nbytes_shared_total, nbatch_fa, true, true, true, use_sparse, warp_size_host, async_kv_preload);
}