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

CUDA : support head_dim 512 with gqa_ratio % 8 (unblocks Gemma 4 12B) (#1921) 

The MMA flash-attention dispatcher only instantiated ncols2 = 8 and 4 for
head_dim 512, so any other GQA ratio hit GGML_ABORT. Gemma 4 12B's global
attention layers use head_dim 512 with a 16:1 GQA ratio (16 query heads /
1 KV head), which aborts at load. Because MTP speculative decoding requires
flash attention, this also blocks the Gemma 4 12B MTP drafter entirely.

Instantiating ncols2 = 16 there is not viable: it exceeds the maximum dynamic
shared memory on Ada (cudaFuncSetAttribute returns invalid argument). Instead,
route gqa_ratio % 8 == 0 (covering 8 and 16) through the existing ncols2 = 8
kernel, which already iterates over Q-head groups (iter_z = ceil(gqa_ratio /
ncols2)). gqa_ratio 8 and 4 behavior is unchanged; this mirrors the divisor
dispatch already used for the 576x512 case below.

Verified on RTX 4070 Ti SUPER (Ada, cc 8.9): Gemma 4 12B + MTP drafter now
runs with flash attention; draft acceptance 43-95% by workload, 1.5-2.2x
end-to-end speedup. The 26B-A4B drafter (gqa_ratio 8) is unaffected.

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Chip Bradford 2026-06-04 08:36:10 -07:00 committed by GitHub
parent 007d640098
commit 19dcc1f7d1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -2255,10 +2255,15 @@ void ggml_cuda_flash_attn_ext_mma_new(ggml_backend_cuda_context & ctx, ggml_tens
return;
}
if (Q->ne[0] == 512 && K->ne[0] == 512 && V->ne[0] == 512) {
if (gqa_ratio == 8) {
// head_dim 512: ncols2=16 exceeds the max dynamic shared memory on some GPUs (e.g. Ada,
// where cudaFuncSetAttribute returns invalid argument), so route gqa_ratio % 8 == 0
// (covers 8 and 16) through the ncols2=8 kernel. It iterates over Q-head groups
// (iter_z = ceil(gqa_ratio/ncols2)), so 16 heads run as two passes of 8. This unblocks
// head_dim-512 models with a 16:1 GQA ratio such as Gemma 4 12B's global layers.
if (gqa_ratio % 8 == 0) {
ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<512, 512, 8>(ctx, dst);
}
else if (gqa_ratio == 4) {
else if (gqa_ratio % 4 == 0) {
ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<512, 512, 4>(ctx, dst);
}
else {