Repository: ggml-org/llama.cpp
Pull request: CUDA: use mma FA kernel for gqa > 4 on RTX 4000

For models such as Qwen 2.5 3b with 8 Q heads per K/V head it seems to be better to use the mma FlashAttention kernel than the vector kernel:

| GPU      | Model         |   Microbatch size | Test         |   t/s master |   t/s 069d410b6 |   Speedup |
|:---------|:--------------|------------------:|:-------------|-------------:|----------------:|----------:|
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128        |       287.82 |          287.43 |      1.00 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d4096  |       252.76 |          253.17 |      1.00 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d8192  |       230.98 |          228.35 |      0.99 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d12288 |       212.58 |          211.96 |      1.00 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d16384 |       196.78 |          196.28 |      1.00 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d20480 |       183.69 |          183.53 |      1.00 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d24576 |       171.84 |          172.22 |      1.00 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d28672 |       161.60 |          161.61 |      1.00 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d32768 |       151.92 |          148.61 |      0.98 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d36864 |       143.42 |          144.39 |      1.01 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d40960 |       135.77 |          139.50 |      1.03 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d45056 |       128.66 |          133.57 |      1.04 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d49152 |       121.80 |          127.93 |      1.05 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d53248 |       115.74 |          122.90 |      1.06 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d57344 |       109.87 |          118.40 |      1.08 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d61440 |       104.90 |          113.65 |      1.08 |
| RTX 4090 | qwen2 3B Q4_0 |               512 | tg128@d65536 |       100.22 |          107.77 |      1.08 |

Notably if the GPU is frequency limited this difference is even larger and up to +25%. With a frequency limit it would also be better to use the mma kernel in other cases such as with LLaMA which has 4 Q heads per K/V head. And since datacenter GPUs have lower frequencies than consumer GPUs this implies that choosing kernels solely based on compute capability is suboptimal. It's currently unclear to me how to best retrieve the GPU clocks in a way that considers user-defined limits, I opened a thread in the NVIDIA developer forums and will make a PR once I get a reply.

Changed files:
filename: ggml/src/ggml-cuda/fattn.cu
patch: @@ -315,8 +315,9 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
 
     const bool gqa_opt_applies = ((Q->ne[2] / K->ne[2]) % 2 == 0) && mask; // The mma-based kernels have GQA-specific optimizations
     const bool mma_needs_data_conversion = K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16;
-    const bool mma_faster_for_bs1 = new_mma_available(cc) && gqa_opt_applies &&
-        (Q->ne[3] > 1 || cc < GGML_CUDA_CC_ADA_LOVELACE) && !mma_needs_data_conversion;
+    const bool mma_faster_for_rtx4000 = Q->ne[3] > 1 || (Q->ne[2] > 4*K->ne[2] && K->ne[1] >= 8192);
+    const bool mma_faster_for_bs1 = new_mma_available(cc) && gqa_opt_applies && !mma_needs_data_conversion &&
+        (cc < GGML_CUDA_CC_ADA_LOVELACE || mma_faster_for_rtx4000);
     const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % (2*warp_size) == 0;
     if (Q->ne[1] == 1 && can_use_vector_kernel && !mma_faster_for_bs1) {
         if (prec == GGML_PREC_DEFAULT) {