19dcc1f7d1
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> |
||
|---|---|---|
| .. | ||
| cmake | ||
| include | ||
| src | ||
| .gitignore | ||
| CMakeLists.txt | ||