Repository: ggml-org/llama.cpp
Pull request: vulkan: improve topk perf for large k, fix overflow in unit tests

The top_k shader reduces workgroupsize elements to k elements in each iteration. When k is large, choose a larger workgroup size so it converges in fewer iterations. This helps with the test cases added here: https://github.com/ggml-org/llama.cpp/pull/17004/files#diff-2749fdb8974ec96afa18444a9d546409318b0a862709139b677eee468c479578R8045

I also saw that for very small tensors, test-backend-ops has an overflow that makes it compute incorrect results. This change does the std::min operation as 64b before converting to 32b.

```
before
  TOP_K(type=f32,ne=[2,1,1,1],k=1):                 -1431647801 runs -    -0.00 us/run -        0 kB/run - -382276.15 GB/s
  TOP_K(type=f32,ne=[1,1,1,1],k=1):                    23544 runs -    42.48 us/run -        0 kB/run -    0.00 GB/s
...
  TOP_K(type=f32,ne=[400,1,1,1],k=400):               548864 runs -     1.83 us/run -        3 kB/run -    1.63 GB/s
  TOP_K(type=f32,ne=[1000,1,1,1],k=400):              311296 runs -     3.22 us/run -        5 kB/run -    1.62 GB/s
  TOP_K(type=f32,ne=[65000,1,1,1],k=400):              16384 runs -    63.84 us/run -      255 kB/run -    3.82 GB/s
  TOP_K(type=f32,ne=[200000,1,1,1],k=400):             16384 runs -    83.12 us/run -      782 kB/run -    8.98 GB/s
  TOP_K(type=f32,ne=[400,16,1,1],k=400):              516096 runs -     1.95 us/run -       50 kB/run -   24.45 GB/s
  TOP_K(type=f32,ne=[1000,16,1,1],k=400):             303104 runs -     3.33 us/run -       87 kB/run -   25.05 GB/s
  TOP_K(type=f32,ne=[65000,16,1,1],k=400):             16384 runs -    97.10 us/run -     4087 kB/run -   40.15 GB/s
  TOP_K(type=f32,ne=[200000,16,1,1],k=400):             5358 runs -   211.52 us/run -    12525 kB/run -   56.47 GB/s

after
  TOP_K(type=f32,ne=[2,1,1,1],k=1):                   679936 runs -     1.48 us/run -        0 kB/run -    0.01 GB/s
  TOP_K(type=f32,ne=[1,1,1,1],k=1):                   679936 runs -     1.47 us/run -        0 kB/run -    0.01 GB/s
...
  TOP_K(type=f32,ne=[400,1,1,1],k=400):               548864 runs -     1.83 us/run -        3 kB/run -    1.63 GB/s
  TOP_K(type=f32,ne=[1000,1,1,1],k=400):              311296 runs -     3.22 us/run -        5 kB/run -    1.62 GB/s
  TOP_K(type=f32,ne=[65000,1,1,1],k=400):              49152 runs -    23.47 us/run -      255 kB/run -   10.38 GB/s
  TOP_K(type=f32,ne=[200000,1,1,1],k=400):             32768 runs -    31.72 us/run -      782 kB/run -   23.54 GB/s
  TOP_K(type=f32,ne=[400,16,1,1],k=400):              524288 runs -     1.92 us/run -       50 kB/run -   24.83 GB/s
  TOP_K(type=f32,ne=[1000,16,1,1],k=400):             303104 runs -     3.32 us/run -       87 kB/run -   25.17 GB/s
  TOP_K(type=f32,ne=[65000,16,1,1],k=400):             24576 runs -    45.12 us/run -     4087 kB/run -   86.39 GB/s
  TOP_K(type=f32,ne=[200000,16,1,1],k=400):            10716 runs -   108.55 us/run -    12525 kB/run -  110.04 GB/s
```

Changed files:
filename: ggml/src/ggml-vulkan/ggml-vulkan.cpp
patch: @@ -10239,7 +10239,9 @@ static void ggml_vk_topk(ggml_backend_vk_context * ctx, vk_context& subctx, cons
 
         // Prefer going as small as num_topk_pipelines - 3 for perf reasons.
         // But if K is larger, then we need a larger workgroup
-        uint32_t max_pipeline = num_topk_pipelines - 3;
+        uint32_t max_pipeline = num_topk_pipelines - 1;
+        uint32_t preferred_pipeline = std::max(num_topk_pipelines - 3, (uint32_t)log2f(float(k)) + 2);
+        max_pipeline = std::min(preferred_pipeline, max_pipeline);
         uint32_t min_pipeline = (uint32_t)log2f(float(k)) + 1;
         // require full subgroup
         min_pipeline = std::max(min_pipeline, ctx->device->subgroup_size_log2);

filename: tests/test-backend-ops.cpp
patch: @@ -1446,14 +1446,14 @@ struct test_case {
             const uint64_t target_flops_cpu =   8ULL * GFLOP;
             const uint64_t target_flops_gpu = 100ULL * GFLOP;
             uint64_t target_flops = is_cpu ? target_flops_cpu : target_flops_gpu;
-            n_runs = std::min<int>(ggml_graph_size(gf) - ggml_graph_n_nodes(gf), target_flops / op_flops(out)) + 1;
+            n_runs = (int)std::min<int64_t>(ggml_graph_size(gf) - ggml_graph_n_nodes(gf), target_flops / op_flops(out)) + 1;
         } else {
             // based on memory size
             const size_t GB = 1ULL << 30;
             const size_t target_size_cpu =  8 * GB;
             const size_t target_size_gpu = 32 * GB;
             size_t target_size = is_cpu ? target_size_cpu : target_size_gpu;
-            n_runs = std::min<int>(ggml_graph_size(gf) - ggml_graph_n_nodes(gf), target_size / op_size(out)) + 1;
+            n_runs = (int)std::min<int64_t>(ggml_graph_size(gf) - ggml_graph_n_nodes(gf), target_size / op_size(out)) + 1;
         }
 
         // duplicate the op
@@ -8043,7 +8043,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
     }
 
     test_cases.emplace_back(new test_argsort(GGML_TYPE_F32, {65000, 16, 1, 1}));
-    for (auto k : {1, 10, 40}) {
+
+    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2, 1, 1, 1}, 1));
+    for (auto k : {1, 10, 40, 400}) {
         for (auto nrows : {1, 16}) {
             for (auto cols : {k, 1000, 65000, 200000}) {
                 test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, k));