From 39febd2adb947b70da35e95f9351b3bf58ad170a Mon Sep 17 00:00:00 2001 From: Frosty40 Date: Wed, 28 Oct 2026 09:18:07 -0500 Subject: [PATCH] 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 | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 97ad492e5..807736e46 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/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()); + const int winner = sycl::reduce_over_group(grp, !taken && cand == best ? tid : INT_MAX, + sycl::minimum()); + 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);