[sb32445]

Overview

With a GQA ratio of 6 (Bonsai 2 27B and Qwen3.5: 24 Q heads on 4 KV heads) the MMA flash attention dispatcher picks ncols2 = 8, so two of the eight head slots of every MMA tile are padding (25 % of the MMA work of the attention kernel). For prompt batches (Q->ne[1] >= 64) ncols2 = 2 computes without padding, and the repeated K/V loads cost less than the saved MMA work. Decode and MTP verify (few tokens) keep ncols2 = 8; without the token threshold decode lost 12 % at 32k depth and 26 % at 96k, so the threshold matters.

RTX 4070, Bonsai 2 27B PTQ1_0 + MTP head, q4_0 K/V: llama-bench pp2048 (-ub 512) +7.0 % at depth 32768 and +14.7 % at depth 98304, llama-server prefill of 114688 tokens 157.4 -> 143.2 s (-9.0 %). Depth 0 and decode are unchanged.

Additional information

• Seven added lines in fattn.cu (one if before the existing gqa_ratio > 4 branch) plus 54 new FLASH_ATTN_EXT test cases (head size 256, GQA ratio 6, 1 to 512 query tokens, f16/q8_0/q4_0 K/V).
• Only ratio 6 and only Q->ne[1] >= 64 are affected; every other model/shape takes the old path.
• Not bit-identical for large batches (different accumulation split): batches below 64 tokens are identical, large batches change by float rounding. KL divergence against the unpatched kernel is 0.00077 (16x512) and 0.00071 (4096 tokens per batch); perplexity 14.5263 against 14.5316 (error 0.62).
• Decode kernels are not touched. In the MTP decode benchmark the outputs are identical and the speed is unchanged (+0.00 %). At 32k depth the output hash differs from the unpatched build in the depth test (rounding in the prefill changes the KV values), which also changes the draft acceptance (88 vs. 82 %) - that decode difference is not a kernel effect and I do not claim it.
• Tuned and measured on one GPU (RTX 4070, Ada, cc 8.9). The 64-token threshold was chosen from a decode/prefill trade-off on this GPU; other architectures (Hopper, Blackwell, RDNA) may want other values. I can restrict it to a compute capability if you prefer.

Test results

• Hardware / software: RTX 4070 12 GB (AD104, cc 8.9), Linux 6.18, NVIDIA driver 615.71, CUDA 13.4, GCC 16.2; Release build, -DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=89.
• Base: prism at c96baf4ab, this commit only (separate builds of base and patch, build 10827 / 10828).
• Model: Ternary-Bonsai-2-27B (PTQ1_0) with a community MTP draft head (Q8_0), q4_0 K/V cache.
• Prefill (llama-bench -p 2048 -n 0 -ub 512 -d 0,32768,98304 -r 2, two runs per build in the order base, patch, patch, base):

Depth	base ↑	patch ↑	change
0	1363.1	1361.3	-0.1 %
32768	960.7	1027.7	+7.0 %
98304	576.8	661.7	+14.7 %

(tok/s, mean of two runs; run-to-run spread ≤ 0.6 %.)

• Server prefill (llama-server, MTP n-max 2, 114688-token prompt): 157.4 -> 143.2 s (-9.0 %); at 32k 30.2 -> 29.1 s.
• Decode (MTP --spec-draft-n-max 2, 4 greedy prompts x 256 tokens, 8 interleaved A/B pairs, paired difference with 95 % bootstrap interval): 115.97 -> 115.98 tok/s, +0.00 % (CI [-0.04, +0.03] %), outputs identical, acceptance 74.0 % in both.
• Quality: perplexity (German text, 16 x 512): 14.5316 -> 14.5263; KL divergence against base 0.000767 ± 0.000024 (16x512) and 0.000706 ± 0.000015 (2 x 4096 batch), same top-p 98.3 / 98.2 %.
• test-backend-ops test -b CUDA0 -o FLASH_ATTN_EXT: 4401/4401 passed (4347 existing + 54 new); -o MUL_MAT: 1669/1669.
• Not tested: other GPUs, other GQA ratios, other models, several slots, multi-GPU.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: The patches were developed with Claude Code (Anthropic's coding agent): it wrote the code, the measurement scripts and the first drafts of the commit messages and PR texts. I decided what to work on (which kernels to optimise, based on profiles of my own setup). The measurements and checks listed above were run in the Claude Code sessions; I did not re-run them independently. I will maintain the changes. Commits where Claude Code was used carry a Co-Authored-By trailer.

🤖 Generated with Claude Code