nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

CUDA: handle GQA = 12 for head size = 256 via new MMA (#2372) 

* iCUDA: handle GQA = 12 for head size = 256 via new MMA

* Cleanup
This commit is contained in:
Kawrakow 2026-08-28 18:10:50 +02:00 committed by GitHub
parent 7cff686d37
commit 555330fbba
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
ggml/src/ggml-cuda/fattn-new-mma.cu
View File

@ -2277,7 +2277,7 @@ void ggml_cuda_flash_attn_ext_mma_new(ggml_backend_cuda_context & ctx, ggml_tens
}
if (K->ne[0] == 256) {
GGML_ASSERT(Q->ne[0] == 256 && V->ne[0] == 256);
if (gqa_ratio == 6) {
if (gqa_ratio % 6 == 0) {
ggml_cuda_flash_attn_ext_mma_f16_case<256, 256, 1, 8>(ctx, dst);
} else {
GGML_ABORT("Not implemented");

6
ggml/src/ggml-cuda/fattn.cu
View File

@ -117,13 +117,15 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
return;
}
int gqa_ratio = Q->ne[2] / K->ne[2];
if (new_mma_available(cc) && K->ne[0] == 128 && V->ne[0] == 128 && Q->ne[0] == 128 && Q->ne[1] == 1 &&
(Q->ne[2] / K->ne[2] == 12 || Q->ne[2] / K->ne[2] == 6 || Q->ne[2] / K->ne[2] == 10)) {
(gqa_ratio == 12 || gqa_ratio == 6 || gqa_ratio == 10)) {
ggml_cuda_flash_attn_ext_mma_new(ctx, dst);
return;
}
if (new_mma_available(cc) && K->ne[0] == 256 && V->ne[0] == 256 && Q->ne[0] == 256 && Q->ne[1] == 1 && Q->ne[2] / K->ne[2] == 6) {
if (new_mma_available(cc) && K->ne[0] == 256 && V->ne[0] == 256 && Q->ne[0] == 256 && Q->ne[1] == 1 && (gqa_ratio == 6 || gqa_ratio == 12)) {
ggml_cuda_flash_attn_ext_mma_new(ctx, dst);
return;
}