Repository: ggml-org/llama.cpp
Pull request: metal : optimize multi-sequence FA vec kernel

ref #10860 #13488

This should largely resolve the text-generation performance for multiple sequences with large prompts on Metal. I think this practically achieves the same effect as PagedAttention. This PR implements it for the FA-vec kernel - we simply skip fully-masked KV cache blocks with the size of the simdgroup (32). The other FA kernel for BS >= 4 already has a [similar optimizaion](https://github.com/ggml-org/llama.cpp/blob/cf0a43bb6490bd49344775abb22ba26f8047cb54/ggml/src/ggml-metal/ggml-metal.metal#L3378-L3398), ~but I think it is not as optimal as it can be. It can now be improved by utilizing #12850 to precompute the skip conditions and pass those to the FA kernel as an extra tensor - this will be done in follow-up PR.~ (edit: nvm - it's already good enough).

Note that this also fixes the TG performance for SWA models, without having to do the defrag (see https://github.com/ggml-org/llama.cpp/pull/13194#discussion_r2084837201).

To test this, we can use the `llama-batched-bench` tool like this. It's important to generate more than one large prompt (via the `-npl` argument) in order to simulate a server with multiple slots. We observe that the TG speed is improved at large contexts.

```bash
make -j && ./bin/llama-batched-bench -m ../models/qwen2.5-7b-coder/ggml-model-q4_k.gguf -c 40000 -b 2048 -ub 512 -npp 0,512,4096,8192 -ntg 32 -npl 2,3 -fa
```

- master

main: n_kv_max = 40192, n_batch = 2048, n_ubatch = 512, flash_attn = 1, is_pp_shared = 0, n_gpu_layers = -1, n_threads = 16, n_threads_batch = 16
|    PP |     TG |    B |   N_KV |   T_PP s | S_PP t/s |   T_TG s | S_TG t/s |      T s |    S t/s |
|-------|--------|------|--------|----------|----------|----------|----------|----------|----------|
|     0 |     32 |    2 |     64 |    0.171 |     0.00 |    0.555 |   115.21 |    0.726 |    88.11 |
|     0 |     32 |    3 |     96 |    0.021 |     0.00 |    0.741 |   129.53 |    0.763 |   125.89 |
|   512 |     32 |    2 |   1088 |    0.856 |  1196.64 |    0.599 |   106.76 |    1.455 |   747.67 |
|   512 |     32 |    3 |   1632 |    1.267 |  1212.00 |    0.788 |   121.85 |    2.055 |   794.09 |
|  4096 |     32 |    2 |   8256 |    7.339 |  1116.21 |    0.659 |    97.09 |    7.998 |  1032.22 |
|  4096 |     32 |    3 |  12384 |   11.070 |  1110.04 |    1.063 |    90.31 |   12.133 |  1020.70 |
|  8192 |     32 |    2 |  16448 |   16.144 |  1014.89 |    0.767 |    83.44 |   16.911 |   972.64 |
|  8192 |     32 |    3 |  24672 |   24.644 |   997.26 |    1.408 |    68.20 |   26.051 |   947.06 |

- PR

main: n_kv_max = 40192, n_batch = 2048, n_ubatch = 512, flash_attn = 1, is_pp_shared = 0, n_gpu_layers = -1, n_threads = 16, n_threads_batch = 16
|    PP |     TG |    B |   N_KV |   T_PP s | S_PP t/s |   T_TG s | S_TG t/s |      T s |    S t/s |
|-------|--------|------|--------|----------|----------|----------|----------|----------|----------|
|     0 |     32 |    2 |     64 |    0.178 |     0.00 |    0.555 |   115.35 |    0.733 |    87.36 |
|     0 |     32 |    3 |     96 |    0.021 |     0.00 |    0.741 |   129.50 |    0.762 |   125.95 |
|   512 |     32 |    2 |   1088 |    0.860 |  1191.01 |    0.602 |   106.39 |    1.461 |   744.53 |
|   512 |     32 |    3 |   1632 |    1.270 |  1209.18 |    0.762 |   125.96 |    2.032 |   802.97 |
|  4096 |     32 |    2 |   8256 |    7.317 |  1119.61 |    0.611 |   104.67 |    7.928 |  1041.33 |
|  4096 |     32 |    3 |  12384 |   11.086 |  1108.44 |    0.869 |   110.43 |   11.955 |  1035.86 |
|  8192 |     32 |    2 |  16448 |   16.128 |  1015.88 |    0.670 |    95.49 |   16.798 |   979.16 |
|  8192 |     32 |    3 |  24672 |   24.585 |   999.62 |    1.004 |    95.61 |   25.590 |   964.14 |


Changed files:
filename: ggml/src/ggml-metal/ggml-metal.metal
patch: @@ -3887,6 +3887,11 @@ kernel void kernel_flash_attn_ext_vec(
                 sm[tiisg] = pm[ic + tiisg];
             }
 
+            // skip -INF blocks
+            if (simd_max(sm[tiisg]) == -INFINITY) {
+                continue;
+            }
+
             // Q*K^T
             {
                 // each simdgroup processes 1 query and NE (NW/NL) head elements