CUDA: fix Gemma 2 numerical issues for FA

[JohannesGaessler]

Fixup to #8542 .

There seem to be numerical issues which (depending on the inputs) can cause incorrect results with Gemma 2 when using FlashAttention. This PR sets the FA precision to FP32 which fixes the issue described in #8542 (comment) .

[strawberrymelonpanda]

Just pulled the PR and checked, looks to be solved to me. Nice work!

./llama-cli --model "gemma-2-27b-it-imat-Q6_K (bartowski).gguf" --prompt "<start_of_turn>user\nhow many squares are on a chessboard?<end_of_turn>\n<start_of_turn>model\n" --verbose --special --flash-attn -ngl 99

This is a classic riddle! There are more squares on a chessboard than you might initially think. Here's how to figure it out: 
[...]

Also checked another Gemma 2 model (gemma-2-27b-it-SimPO-37K-Q5_K_M.gguf) and Gemma 2 9B, both fine.

[strawberrymelonpanda]

Benchmarks for what it's worth:

gemma-2-27b-it-imat-Q6_K_(bartowski).gguf

Device 0: NVIDIA GeForce RTX 3090, compute capability 8.6, VMM: yes
| model                          |       size |     params | backend    | ngl | fa |          test |              t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | -: | ------------: | ---------------: |
| gemma2 27B Q6_K                |  21.70 GiB |    28.41 B | CUDA       |  99 |  0 |         pp512 |  1149.13 ± 15.39 |
| gemma2 27B Q6_K                |  21.70 GiB |    28.41 B | CUDA       |  99 |  0 |         tg128 |     29.53 ± 0.28 |
| gemma2 27B Q6_K                |  21.70 GiB |    28.41 B | CUDA       |  99 |  1 |         pp512 |  1230.19 ± 10.37 |
| gemma2 27B Q6_K                |  21.70 GiB |    28.41 B | CUDA       |  99 |  1 |         tg128 |     30.58 ± 0.12 |

gemma-2-27b-it-SimPO-37K-Q5_K_M.gguf

Device 0: NVIDIA GeForce RTX 3090, compute capability 8.6, VMM: yes
| model                          |       size |     params | backend    | ngl | fa |          test |              t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | -: | ------------: | ---------------: |
| gemma2 27B Q5_K - Medium       |  18.97 GiB |    28.41 B | CUDA       |  99 |  0 |         pp512 |   1299.05 ± 8.52 |
| gemma2 27B Q5_K - Medium       |  18.97 GiB |    28.41 B | CUDA       |  99 |  0 |         tg128 |     34.25 ± 0.23 |
| gemma2 27B Q5_K - Medium       |  18.97 GiB |    28.41 B | CUDA       |  99 |  1 |         pp512 |  1350.98 ± 22.11 |
| gemma2 27B Q5_K - Medium       |  18.97 GiB |    28.41 B | CUDA       |  99 |  1 |         tg128 |     35.62 ± 0.23 |

[JohannesGaessler]

You should be seeing more of a difference if you set higher values for -n and -p since with a mostly empty context the attention is comparatively fast.

[strawberrymelonpanda]

Thanks for the tip. Mostly just testing and wanted to show setting FP32 doesn't seem to cause any issues or slowdown.

Here's some higher values.

| model                          |       size |     params | backend    | ngl | fa |          test |              t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | -: | ------------: | ---------------: |
| gemma2 27B Q5_K - Medium       |  18.97 GiB |    28.41 B | CUDA       |  99 |  0 |        pp4096 |   1052.16 ± 8.75 |
| gemma2 27B Q5_K - Medium       |  18.97 GiB |    28.41 B | CUDA       |  99 |  0 |        tg1024 |     33.08 ± 0.09 |
| gemma2 27B Q5_K - Medium       |  18.97 GiB |    28.41 B | CUDA       |  99 |  1 |        pp4096 |   1256.27 ± 5.57 |
| gemma2 27B Q5_K - Medium       |  18.97 GiB |    28.41 B | CUDA       |  99 |  1 |        tg1024 |     35.13 ± 0.17 |

Looks like a nice improvement to me.