Repository: ggml-org/llama.cpp
Pull request: metal : make the FA extra sizes consistent

While looking into https://github.com/ggml-org/llama.cpp/issues/17033#issuecomment-3508789138 found that the warmup in `llama-batched-bench` trashes the worst-case graph allocation using the Metal backend, causing extra graph allocations later on. The reason is because the [extra FA fleeting memory](https://github.com/ggml-org/llama.cpp/blob/1032256ec95986f60124a62ace6a628106546497/ggml/src/ggml-metal/ggml-metal.cpp#L194-L200) for the [small warmup batch](https://github.com/ggml-org/llama.cpp/blob/1032256ec95986f60124a62ace6a628106546497/tools/batched-bench/batched-bench.cpp#L102-L112) ends up being larger than the memory for the [worst case estimate](https://github.com/ggml-org/llama.cpp/blob/1032256ec95986f60124a62ace6a628106546497/src/llama-context.cpp#L387-L399).

Fix by making the extra size more correlated with the input shapes.

```bash
make -j && ./bin/llama-batched-bench -m ../models/qwen2.5-3b-coder/ggml-model-q8_0.gguf -c 150792 -npp 8192 -ntg 32 -npl 1,2,4,8,16 -kvu -tgs --no-mmap
```

### master

main: n_kv_max = 151040, n_batch = 2048, n_ubatch = 512, flash_attn = -1, is_pp_shared = 0, is_tg_separate = 1, 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 |
|-------|--------|------|--------|----------|----------|----------|----------|----------|----------|
|  8192 |     32 |    1 |   8224 |    3.406 |  2405.30 |    0.289 |   110.58 |    3.695 |  2225.58 |
|  8192 |     32 |    2 |  16448 |    6.843 |  2394.42 |    0.590 |   108.46 |    7.433 |  2212.94 |
|  8192 |     32 |    4 |  32896 |   14.106 |  2322.90 |    1.238 |   103.37 |   15.345 |  2143.80 |
|  8192 |     32 |    8 |  65792 |   29.769 |  2201.46 |    2.767 |    92.53 |   32.536 |  2022.12 |
|  8192 |     32 |   16 | 131584 |   65.512 |  2000.74 |    6.663 |    76.84 |   72.175 |  1823.13 |

### PR

main: n_kv_max = 151040, n_batch = 2048, n_ubatch = 512, flash_attn = -1, is_pp_shared = 0, is_tg_separate = 1, 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 |
|-------|--------|------|--------|----------|----------|----------|----------|----------|----------|
|  8192 |     32 |    1 |   8224 |    3.287 |  2491.90 |    0.289 |   110.62 |    3.577 |  2299.30 |
|  8192 |     32 |    2 |  16448 |    6.646 |  2465.21 |    0.589 |   108.59 |    7.235 |  2273.25 |
|  8192 |     32 |    4 |  32896 |   13.581 |  2412.81 |    1.240 |   103.19 |   14.821 |  2219.51 |
|  8192 |     32 |    8 |  65792 |   28.191 |  2324.72 |    2.751 |    93.04 |   30.942 |  2126.28 |
|  8192 |     32 |   16 | 131584 |   60.293 |  2173.91 |    6.633 |    77.19 |   66.927 |  1966.09 |

Changed files:
filename: ggml/src/ggml-metal/ggml-metal-ops.cpp
patch: @@ -1975,7 +1975,9 @@ size_t ggml_metal_op_flash_attn_ext_extra_pad(const ggml_tensor * op) {
     const bool has_mask = op->src[3] != nullptr;
 
     if (ggml_metal_op_flash_attn_ext_use_vec(op)) {
-        const bool has_kvpad = ne11 % OP_FLASH_ATTN_EXT_VEC_NCPSG != 0;
+        // note: always reserve the padding space to avoid graph reallocations
+        //const bool has_kvpad = ne11 % OP_FLASH_ATTN_EXT_VEC_NCPSG != 0;
+        const bool has_kvpad = true;
 
         if (has_kvpad) {
             res += OP_FLASH_ATTN_EXT_VEC_NCPSG*(
@@ -1984,7 +1986,8 @@ size_t ggml_metal_op_flash_attn_ext_extra_pad(const ggml_tensor * op) {
                 (has_mask ? ggml_type_size(GGML_TYPE_F16)*ne31*ne32*ne33 : 0));
         }
     } else {
-        const bool has_kvpad = ne11 % OP_FLASH_ATTN_EXT_NCPSG != 0;
+        //const bool has_kvpad = ne11 % OP_FLASH_ATTN_EXT_NCPSG != 0;
+        const bool has_kvpad = true;
 
         if (has_kvpad) {
             res += OP_FLASH_ATTN_EXT_NCPSG*(
@@ -2020,9 +2023,10 @@ size_t ggml_metal_op_flash_attn_ext_extra_blk(const ggml_tensor * op) {
     const bool is_vec = ggml_metal_op_flash_attn_ext_use_vec(op);
 
     // this optimization is not useful for the vector kernels
-    if (is_vec) {
-        return res;
-    }
+    // note: always reserve the blk buffer to avoid graph reallocations
+    //if (is_vec) {
+    //    return res;
+    //}
 
     const int nqptg = is_vec ? OP_FLASH_ATTN_EXT_VEC_NQPTG : OP_FLASH_ATTN_EXT_NQPTG;
     const int ncpsg = is_vec ? OP_FLASH_ATTN_EXT_VEC_NCPSG : OP_FLASH_ATTN_EXT_NCPSG;
@@ -2049,13 +2053,16 @@ size_t ggml_metal_op_flash_attn_ext_extra_tmp(const ggml_tensor * op) {
 
     size_t res = 0;
 
-    if (ggml_metal_op_flash_attn_ext_use_vec(op)) {
+    // note: always reserve the temp buffer to avoid graph reallocations
+    //if (ggml_metal_op_flash_attn_ext_use_vec(op)) {
+    if (true) {
         const int64_t nwg = 32;
+        const int64_t ne01_max = std::min(ne01, 32);
 
         // temp buffer for writing the results from each workgroup
         // - ne20: the size of the Value head
         // -  + 2: the S and M values for each intermediate result
-        res += ggml_type_size(GGML_TYPE_F32)*(ne01*ne02*ne03*nwg*(ne20 + 2));
+        res += ggml_type_size(GGML_TYPE_F32)*(ne01_max*ne02*ne03*nwg*(ne20 + 2));
     }
 
     return res;