CUDA FA: faster TG when GQA is 16 and head size is 128

This commit is contained in:
Kawrakow 2026-06-15 11:46:02 +00:00
parent f81673c7db
commit 6be3a488d3
4 changed files with 19 additions and 0 deletions

View File

@ -223,6 +223,19 @@ void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, ggml_tens
return;
}
if (use_gqa_opt && gqa_ratio % 16 == 0 && Q->ne[0] == 128 && Q->ne[0] <= 8) {
if (Q->ne[1] <= 1) {
ggml_cuda_flash_attn_ext_mma_f16_case<128, 1, 16>(ctx, dst);
}
else if (Q->ne[1] <= 2) {
ggml_cuda_flash_attn_ext_mma_f16_case<128, 2, 16>(ctx, dst);
}
else {
ggml_cuda_flash_attn_ext_mma_f16_case<128, 4, 16>(ctx, dst);
}
return;
}
if (use_gqa_opt && gqa_ratio % 8 == 0) {
ggml_cuda_flash_attn_ext_mma_f16_switch_hs<8>(ctx, dst);
return;

View File

@ -8,3 +8,5 @@ DECL_FATTN_MMA_F16_CASE(96, 1, 8);
DECL_FATTN_MMA_F16_CASE(112, 1, 8);
DECL_FATTN_MMA_F16_CASE(128, 1, 8);
DECL_FATTN_MMA_F16_CASE(256, 1, 8);
DECL_FATTN_MMA_F16_CASE(128, 1, 16);

View File

@ -8,3 +8,5 @@ DECL_FATTN_MMA_F16_CASE(96, 2, 8);
DECL_FATTN_MMA_F16_CASE(112, 2, 8);
DECL_FATTN_MMA_F16_CASE(128, 2, 8);
DECL_FATTN_MMA_F16_CASE(256, 2, 8);
DECL_FATTN_MMA_F16_CASE(128, 2, 16);

View File

@ -8,3 +8,5 @@ DECL_FATTN_MMA_F16_CASE(96, 4, 8);
DECL_FATTN_MMA_F16_CASE(112, 4, 8);
DECL_FATTN_MMA_F16_CASE(128, 4, 8);
DECL_FATTN_MMA_F16_CASE(256, 4, 8);
DECL_FATTN_MMA_F16_CASE(128, 4, 16);