Repository: ggml-org/llama.cpp
Pull request: (Bugfix, ggml-cuda) Pool alloc count fix + small size computation type adjustment

TLDR: One of sizes passed to the pool alloc was count of bytes instead of count of objects in fattn-common::launch_fattn, added explicit cast of size variable type before multiplication into argsort.cu::argsort_f32_i32_cuda_cub, builtin tests (`./test-backend-ops test` command) and inference experiments didn't reveal anomalies after the change.

<details>
<summary>System setup</summary>

Ubuntu 24.04, amd ry7 5800x3d, b550 chipset
```
ggml_cuda_init: found 3 CUDA devices:
  Device 0: NVIDIA GeForce RTX 3090, compute capability 8.6, VMM: yes
  Device 1: Tesla V100-SXM2-16GB, compute capability 7.0, VMM: yes
  Device 2: Tesla V100-SXM2-16GB, compute capability 7.0, VMM: yes
main: n_parallel is set to auto, using n_parallel = 4 and kv_unified = true
build: 7610 (c6f0e832d) with GNU 14.2.0 for Linux x86_64
system info: n_threads = 8, n_threads_batch = 8, total_threads = 16

system_info: n_threads = 8 (n_threads_batch = 8) / 16 | CUDA : ARCHS = 700,860 | USE_GRAPHS = 1 | PEER_MAX_BATCH_SIZE = 128 | CPU : SSE3 = 1 | SSSE3 = 1 | AVX = 1 | AVX2 = 1 | F16C = 1 | FMA = 1 | BMI2 = 1 | LLAMAFILE = 1 | OPENMP = 1 | REPACK = 1 |
```
don't ask what have I done to fit 2 tesla v100 with rtx 3090 onto ordinary desktop board :)

</details>

AI usage disclosure and description of "how did I find it":
I have had a suspicion that there was a memory leak, as I was facing spontaneous oom errors with very similar description mentioning vmm related function, even when leaving around 800mb to spare at least on each of the gpus, so I gave AI agent with IDE harness detailed description and logs for it to investigate, but as it turned out it was user error (I have underestimated memory consumption of large context window + default implicit --parallel 4).
However it pointed out that one of the pool allocating functions was passed with byte count instead of object count, resulting in 8x over-allocation. I of course looked myself and made sure that the math was mathing (the multiplication factor of 2*sizeof(float) was indeed unintended) and even asked it to look around and double check that we both are right. Also I have asked it to find other size math issues, though it only pointed out the int multiplied by int hypothetical overflow in argsort (I doubt it can cause any troubles, but decided that explicit cast won't hurt) after looking through the codebase.

Then I have MANUALLY adjusted the code, built the solution, ran `./test-backend-ops test` for all devices and cpu, tested some of my commonly used llms for anomalies or crashes, performed some of the runs with `CUDA_LAUNCH_BLOCKING=1 compute-sanitizer --tool memcheck <llama bin> <args>`, monitored nvtop for weird dynamics or mismatches with memory_breakdown_print.

AI wasn't used for performing actual coding, making commits or writing description text for PR, only for codebase sifting and read-only investigation of the issue



Changed files:
filename: ggml/src/ggml-cuda/argsort.cu
patch: @@ -29,8 +29,8 @@ static void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
                                      const int        nrows,
                                      ggml_sort_order  order,
                                      cudaStream_t     stream) {
-    ggml_cuda_pool_alloc<int>   temp_indices_alloc(pool, ncols * nrows);
-    ggml_cuda_pool_alloc<float> temp_keys_alloc(pool, ncols * nrows);
+    ggml_cuda_pool_alloc<int>   temp_indices_alloc(pool, ((size_t) ncols) * nrows);
+    ggml_cuda_pool_alloc<float> temp_keys_alloc(pool, ((size_t) ncols) * nrows);
     ggml_cuda_pool_alloc<int>   offsets_alloc(pool, nrows + 1);
 
     int *   temp_indices = temp_indices_alloc.get();

filename: ggml/src/ggml-cuda/fattn-common.cuh
patch: @@ -918,7 +918,7 @@ void launch_fattn(
         blocks_num.y = 1;
         blocks_num.z = 1;
 
-        dst_tmp_meta.alloc(blocks_num.x*ncols * (2*2 + DV) * sizeof(float));
+        dst_tmp_meta.alloc(((size_t) blocks_num.x) * ncols * (2 + DV/2));
     } else {
         const int ntiles_KQ = (K->ne[1] + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by tensor size.