sycl: work-group argmax for short TOP_K rows (#31110)

* sycl: work-group argmax for short TOP_K rows

Used angelX to optimize SYCL/drama 3 coordination, results are as follows

Rows with ncols <= max work-group size: one work-item per column, k rounds of work-group argmax. Ties take the lower column, NaN is never picked. Longer rows keep the existing paths.

Arc Pro B70, oneAPI 2026.1, test-backend-ops perf (existing short-row cases), base 9b56ccf00 vs patch, median ratio of 6 rotating A/B/A rounds:

| ncols | nrows | k | base us | patch us | speedup |
|---|---|---|---|---|---|
| 16 | 1024 | 1 | 19.97 | 2.88 | 6.95x |
| 16 | 1024 | 2 | 23.83 | 2.90 | 8.25x |
| 16 | 1024 | 8 | 66.12 | 4.66 | 14.21x |
| 16 | 8192 | 1 | 110.42 | 5.96 | 18.51x |
| 16 | 8192 | 2 | 130.21 | 8.82 | 14.77x |
| 16 | 8192 | 8 | 420.62 | 27.27 | 15.42x |
| 128 | 1024 | 1 | 20.49 | 2.94 | 6.97x |
| 128 | 1024 | 2 | 25.89 | 3.45 | 7.52x |
| 128 | 1024 | 8 | 101.67 | 8.95 | 11.36x |
| 128 | 8192 | 1 | 112.19 | 13.29 | 8.44x |
| 128 | 8192 | 2 | 141.18 | 20.60 | 6.85x |
| 128 | 8192 | 8 | 643.21 | 67.60 | 9.51x |
| 1024 | 1024 | 1 | 17.32 | 14.98 | 1.16x |
| 1024 | 1024 | 2 | 28.77 | 24.02 | 1.20x |
| 1024 | 1024 | 8 | 151.87 | 77.88 | 1.95x |
| 1024 | 8192 | 1 | 156.49 | 138.78 | 1.13x |
| 1024 | 8192 | 2 | 201.47 | 209.18 | 0.96x |
| 1024 | 8192 | 8 | 1006.25 | 632.60 | 1.59x |

The 1024 x 8192, k=2 case is 3.7% slower.

Correctness: test-backend-ops -o TOP_K 533/533 and -o ARGSORT 98/98 on the patch, again on the rebase onto 3f03b1772.

Full test-backend-ops aborts on base and patch alike at the non-F32 SWIGLU_OAI assert (element_wise.cpp:1187, see #30857). With only those cases excluded: 20145/20173 on base and patch. One CPY f16->q2_0 [256,2,3,4] failure was seen once on the patch, and 3/3 re-runs pass on base and patch. All other patch-run failures also fail on base.

Assisted-by: Claude Opus 5.5

* made changes as suggested

Assisted-by: Claude Opus 5.5

ggml/src/ggml-sycl/ggml-sycl.cpp

@@ -3010,6 +3010,35 @@ static void top_k_f32_sycl(
         return;
     }
 
+    // Short rows: one work-item per column and k work-group argmax rounds.
+    // Ties take the lower column.
+    const int wg = std::max(WARP_SIZE, next_power_of_2((int) ncols));
+    if (wg <= ggml_sycl_info().max_work_group_sizes[ctx.device]) {
+        main_stream->parallel_for(
+            sycl::nd_range<1>(sycl::range<1>(nrows * wg), sycl::range<1>(wg)),
+            [=](sycl::nd_item<1> item_ct1) {
+                const int64_t row = item_ct1.get_group(0);
+                const int     tid = item_ct1.get_local_id(0);
+                const auto    grp = item_ct1.get_group();
+
+                // NaN is never preferred, and every round still has a winner since k <= ncols
+                const float x     = tid < ncols ? src[row * ncols + tid] : -INFINITY;
+                const float val   = sycl::isnan(x) ? -INFINITY : x;
+                bool        taken = tid >= ncols;
+                for (int i = 0; i < k; i++) {
+                    const float cand   = taken ? -INFINITY : val;
+                    const float best   = sycl::reduce_over_group(grp, cand, sycl::maximum<float>());
+                    const int   winner = sycl::reduce_over_group(grp, !taken && cand == best ? tid : INT_MAX,
+                                                                 sycl::minimum<int>());
+                    if (tid == winner) {
+                        dst_indices[row * k + i] = tid;
+                        taken = true;
+                    }
+                }
+            });
+        return;
+    }
+
     const int block_size = 128;
 
     const sycl::range<1> block_dims(block_size);