From 19dcc1f7d122460dcc77c9b1a9cf02bfc658b91e Mon Sep 17 00:00:00 2001 From: Chip Bradford Date: Thu, 4 Jun 2026 08:36:10 -0700 Subject: [PATCH] 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) --- ggml/src/ggml-cuda/fattn-new-mma.cu | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-new-mma.cu b/ggml/src/ggml-cuda/fattn-new-mma.cu index b2f67e4e..8c60b563 100644 --- a/ggml/src/ggml-cuda/fattn-new-mma.cu +++ b/ggml/src/ggml-cuda/fattn-new-mma.cu @@ -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 {