sycl: rfc: Use radix select for top_k (#28670)

* sycl: GPU-resident TOP_K for large k, parallelised over the device

The SYCL backend refused GGML_OP_TOP_K above k = 32 and let it fall back to
the CPU, a backend round-trip per call. The limit was not conservatism: the
scan-merge kernels keep (split_block + 1) * k candidate (value, index) pairs
in SLM, so at k = 128 a work-group already needs 132 KB and cannot launch.
qwen4exp's sparse-attention indexer asks for k = 2048 in 12 layers on every
token, so this fired at every context length.

Add a radix select for large k. The k-th largest is found by four
most-significant-first passes over an order-preserving unsigned key: histogram
the digit over the candidate set, walk the buckets from the top, and recurse
into the one where the running count reaches what is still needed. SLM holds
the histogram rather than candidates, so the footprint is independent of k.
A final pass emits every column beating the pivot plus exactly as many
pivot-equal columns as are still missing, so duplicate keys still yield
exactly k distinct indices. Output order is not required and is not paid for:
ggml-cpu/ops.cpp swaps its first two outputs to say so.

The key folds -0.0 onto +0.0 so its equivalence classes match the reference
comparator, under which the two tie. NaN has no defined order in the reference
(its comparator is not a strict weak order there); here +NaN keys above +inf
and -NaN below -inf, which at least makes the result deterministic.

One work-group per row leaves the device idle whenever a graph has fewer rows
than it has cores, which at batch size 1 means one work-group full stop:
qwen4exp tops-k a tensor of shape [n_kv, n_tokens/n_stream, n_stream], so
token generation gives nrows == 1, and the backend sampler reshapes logits to
a single row as well. Measured, ne=[200000,1] and ne=[200000,16] cost 358.0 us
and 363.4 us -- sixteen rows for 1.5% more wall-clock.

So also spread a row over several groups when there are too few rows to cover
the device. Per-pass state moves to global memory and each digit pass becomes
its own launch, since a work-group barrier can no longer span the row. Groups
accumulate in SLM and contribute 256 global atomics each, keeping global
traffic per-group rather than per-element, and the last group of a row -- the
one whose fetch_add returns G-1 -- performs that pass's scan, holding the
launch count at one per digit plus one emit. The group count comes from the
device and is floor-divided by nrows, so a row count that already covers the
device is left whole and pays nothing. Below 64K columns the single-group
kernel finishes inside the cost of the extra launches and stays in charge.

Reading the row's prefix/mask/need through a device-scope atomic_ref costs
more than the sweep it guards: those loads are uncached, so passes 2-4 ran at
49 us against 12 us for pass 1. One lane reads them into SLM and the group
takes them from there -- 208 us -> 44.6 us at ne=[131072,1], k=2048.

The block size now takes the device's max_work_group_size instead of a cap of
512. The cap was never a floor, so a device reporting 512 is unaffected; one
allowing 1024 was being given half its width.

Finally, put the scan-merge gate where the two paths actually cross. That
kernel's cost climbs with k while the radix select's does not; measured over
widths from 2 to 200K columns and row counts from 1 to 8192, radix is ahead
everywhere from k = 8 up and behind at k <= 2, where scan-merge's smaller
fixed cost wins. The short-row corner (ncols=2, nrows=65536, as in bailingmoe2
group selection) is exactly where radix loses at low k, and the gate keeps it
on scan-merge.

Op-level against the CPU-fallback path this replaces, and against the
single-group radix select for the split: 4.98x at ne=[131072,1] k=2048,
6.65x at ne=[151936,1] k=40, 13.35x at k=20, 118x at ne=[65000,16] k=32.
No measured shape regressed. End to end on 3x Arc Pro B60 with
Qwen3.8-Flash-Next UD-IQ4_XS, llama-bench tg64, the parallelisation is worth
5.91 -> 6.05 t/s at d=131072 and a wash at shallower depths. Perplexity over
wikitext-2 is unchanged within noise at both 512 and 81920 context.

test-backend-ops: 525/525 TOP_K (previously every k > 32 case was refused),
880/880 MUL_MAT_ID. Perf coverage added for k > 32 at large widths and for the
short-row corner, neither of which was exercised before.

* move topk-select to topk-radix.{cpp|hpp}

---------

Co-authored-by: cwriter <cwriter@localhost>

ggml/src/ggml-sycl/backend.hpp

@@ -44,6 +44,7 @@
 #include "ssm_conv.hpp"
 #include "softmax.hpp"
 #include "topk-moe.hpp"
+#include "topk-radix.hpp"
 #include "tsembd.hpp"
 #include "upscale.hpp"
 #include "wkv.hpp"

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

@@ -3121,10 +3121,14 @@ static void ggml_sycl_op_top_k(ggml_backend_sycl_context & ctx, ggml_tensor * ds
     const int64_t ncols = src0->ne[0];
     const int64_t nrows = ggml_nrows(src0);
 
-    GGML_ASSERT(k > 0 && k <= 32);
+    GGML_ASSERT(k > 0);
     GGML_ASSERT(k <= ncols);
 
-    top_k_f32_sycl(ctx, src0_dd, dst_dd, ncols, nrows, k, main_stream);
+    if (k <= SYCL_TOP_K_SCAN_MERGE_MAX_K) {
+        top_k_f32_sycl(ctx, src0_dd, dst_dd, ncols, nrows, k, main_stream);
+    } else {
+        ggml_sycl_top_k_radix(ctx, src0_dd, dst_dd, ncols, nrows, k, main_stream);
+    }
 }
 
 inline void ggml_sycl_op_argmax(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
@@ -6644,7 +6648,7 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
                 op->type == GGML_TYPE_I32 &&
                 src0->type == GGML_TYPE_F32 &&
                 ggml_is_contiguous(src0) &&
-                k > 0 && k <= 32;
+                k > 0 && k <= src0->ne[0];
         }
         case GGML_OP_POOL_2D:
         case GGML_OP_POOL_1D:

ggml/src/ggml-sycl/topk-radix.cpp

@@ -0,0 +1,531 @@
+#include "topk-radix.hpp"
+
+#include "common.hpp"
+
+#include <algorithm>
+
+// Large-k top-k by radix select on an order-preserving unsigned key.
+//
+// The k-th largest key of a row is found by four most-significant-first passes over its
+// 8-bit digits: histogram the digit over the candidate set, walk the buckets from the
+// top, and recurse into the bucket where the running count reaches what is still
+// needed. Everything strictly above that bucket is in the top-k. A final pass emits
+// every column whose key beats the pivot, then exactly as many pivot-equal columns as
+// are still missing, so duplicate keys yield exactly k distinct indices.
+//
+// SLM holds only the histogram, so unlike the scan-merge kernels the cost does not grow
+// with k. One work-group owns a row and runs every pass, so a top-k is one launch and
+// needs no pool scratch. The row is re-read once per pass rather than compacted, which
+// keeps the candidate set implicit: (key & mask) == prefix.
+//
+// The output is the set of winning indices in no particular order, which is what the
+// reference op provides (it swaps its first two outputs to say so) and what
+// test-backend-ops compares.
+
+static constexpr int SYCL_TOP_K_RADIX_BITS       = 8;
+static constexpr int SYCL_TOP_K_RADIX_BUCKETS    = 1 << SYCL_TOP_K_RADIX_BITS;
+// Private histogram copies, interleaved per bucket so neighbouring lanes hit
+// neighbouring banks. Lanes of one instruction spread over the copies, which is what
+// bounds the atomic serialisation on tie-heavy rows.
+static constexpr int SYCL_TOP_K_RADIX_HIST_COPIES = 8;
+static constexpr int SYCL_TOP_K_RADIX_HIST_SIZE   = SYCL_TOP_K_RADIX_BUCKETS * SYCL_TOP_K_RADIX_HIST_COPIES;
+// Past the histogram: pivot digit, pivot bucket count, remaining need, then the two
+// emit counters.
+static constexpr int SYCL_TOP_K_RADIX_SLM_WORDS   = SYCL_TOP_K_RADIX_HIST_SIZE + 5;
+
+// Larger float <=> larger key. The reference comparator is a plain float '>', under which
+// -0.0 and +0.0 tie, so -0.0 is folded onto +0.0 first. NaN has no defined order in the
+// reference (its comparator is not a strict weak order on NaN); here a positive NaN keys
+// above +inf and a negative NaN below -inf, which at least makes the result deterministic.
+static inline uint32_t top_k_radix_key(float f) {
+    uint32_t u = sycl::bit_cast<uint32_t>(f);
+    if (u == 0x80000000u) {
+        u = 0u;
+    }
+    return (u & 0x80000000u) ? ~u : (u | 0x80000000u);
+}
+
+static void top_k_radix_select_f32(
+    const float *   src,
+    int32_t *       dst_idx,
+    const int       ncols,
+    const int       k,
+    uint32_t *      slm,
+    const sycl::nd_item<1> & item_ct1
+) {
+    using local_atomic = sycl::atomic_ref<uint32_t, sycl::memory_order::relaxed,
+                                          sycl::memory_scope::work_group,
+                                          sycl::access::address_space::local_space>;
+
+    const int tid        = item_ct1.get_local_id(0);
+    const int block_size = item_ct1.get_local_range(0);
+
+    uint32_t * hist     = slm;
+    uint32_t * s_digit  = slm + SYCL_TOP_K_RADIX_HIST_SIZE;
+    uint32_t * s_bucket = slm + SYCL_TOP_K_RADIX_HIST_SIZE + 1;
+    uint32_t * s_need   = slm + SYCL_TOP_K_RADIX_HIST_SIZE + 2;
+    uint32_t * s_cnt_gt = slm + SYCL_TOP_K_RADIX_HIST_SIZE + 3;
+    uint32_t * s_cnt_eq = slm + SYCL_TOP_K_RADIX_HIST_SIZE + 4;
+
+    if (tid == 0) {
+        *s_cnt_gt = 0;
+        *s_cnt_eq = 0;
+    }
+
+    const int copy = tid & (SYCL_TOP_K_RADIX_HIST_COPIES - 1);
+
+    uint32_t prefix = 0;   // digits fixed so far, in place
+    uint32_t mask   = 0;   // which bits of prefix are fixed
+    uint32_t need   = (uint32_t) k;
+
+    for (int shift = 32 - SYCL_TOP_K_RADIX_BITS; shift >= 0; shift -= SYCL_TOP_K_RADIX_BITS) {
+        for (int i = tid; i < SYCL_TOP_K_RADIX_HIST_SIZE; i += block_size) {
+            hist[i] = 0;
+        }
+        item_ct1.barrier(sycl::access::fence_space::local_space);
+
+        for (int col = tid; col < ncols; col += block_size) {
+            const uint32_t key = top_k_radix_key(src[col]);
+            if ((key & mask) == prefix) {
+                const uint32_t bucket = (key >> shift) & (SYCL_TOP_K_RADIX_BUCKETS - 1);
+                local_atomic(hist[bucket * SYCL_TOP_K_RADIX_HIST_COPIES + copy]).fetch_add(1u);
+            }
+        }
+        item_ct1.barrier(sycl::access::fence_space::local_space);
+
+        // Lane t takes bucket 255 - t, so an inclusive scan over lanes counts from the top
+        // bucket downward. The pivot is the unique bucket whose cumulative count first
+        // reaches need; the previous cumulative count is what the higher buckets contribute.
+        uint32_t cnt = 0;
+        if (tid < SYCL_TOP_K_RADIX_BUCKETS) {
+            const uint32_t * h = hist + (SYCL_TOP_K_RADIX_BUCKETS - 1 - tid) * SYCL_TOP_K_RADIX_HIST_COPIES;
+            for (int c = 0; c < SYCL_TOP_K_RADIX_HIST_COPIES; c++) {
+                cnt += h[c];
+            }
+        }
+        const uint32_t incl = sycl::inclusive_scan_over_group(item_ct1.get_group(), cnt, sycl::plus<uint32_t>());
+
+        if (tid < SYCL_TOP_K_RADIX_BUCKETS && incl >= need && incl - cnt < need) {
+            *s_digit = (uint32_t) (SYCL_TOP_K_RADIX_BUCKETS - 1 - tid);
+            *s_bucket = cnt;
+            *s_need   = need - (incl - cnt);
+        }
+        item_ct1.barrier(sycl::access::fence_space::local_space);
+
+        const uint32_t digit      = *s_digit;
+        const uint32_t bucket_cnt = *s_bucket;
+        need   = *s_need;
+        prefix |= digit << shift;
+        mask   |= (uint32_t) (SYCL_TOP_K_RADIX_BUCKETS - 1) << shift;
+
+        // Every candidate in the pivot bucket is wanted: the remaining digits cannot
+        // change the answer, and the masked emit below is exact as it stands.
+        if (bucket_cnt == need) {
+            break;
+        }
+        // The next pass rewrites hist and s_*; the reads above must land first.
+        item_ct1.barrier(sycl::access::fence_space::local_space);
+    }
+
+    item_ct1.barrier(sycl::access::fence_space::local_space);
+
+    // Exactly k - need columns have (key & mask) > prefix; the first need of the pivot-equal
+    // columns fill the tail. Both counters live in SLM since the whole row is this group.
+    const uint32_t base_eq = (uint32_t) k - need;
+
+    for (int col = tid; col < ncols; col += block_size) {
+        const uint32_t kp = top_k_radix_key(src[col]) & mask;
+        if (kp > prefix) {
+            const uint32_t pos = local_atomic(*s_cnt_gt).fetch_add(1u);
+            dst_idx[pos] = col;
+        } else if (kp == prefix) {
+            const uint32_t pos = local_atomic(*s_cnt_eq).fetch_add(1u);
+            if (pos < need) {
+                dst_idx[base_eq + pos] = col;
+            }
+        }
+    }
+}
+
+static void top_k_radix_f32_sycl(
+    ggml_backend_sycl_context & ctx,
+    const float * src,
+    int32_t * dst_indices,
+    const int64_t ncols,
+    const int64_t nrows,
+    const int k,
+    dpct::queue_ptr main_stream
+) {
+    GGML_ASSERT(ncols <= INT32_MAX);
+
+    // One group per row; every pass is a strided sweep of the row, so lanes in flight is the
+    // only lever, and the device's own limit is the answer -- there is nothing here that
+    // wants a smaller group. Must still cover the 256 buckets for the scan step.
+    const int block_size = ggml_sycl_info().max_work_group_sizes[ctx.device];
+    GGML_ASSERT(block_size >= SYCL_TOP_K_RADIX_BUCKETS);
+
+    const sycl::range<1> block_dims(block_size);
+    const sycl::range<1> grid_dims(nrows);
+
+    main_stream->submit([&](sycl::handler &cgh) {
+        sycl::local_accessor<uint32_t, 1> slm(sycl::range<1>(SYCL_TOP_K_RADIX_SLM_WORDS), cgh);
+
+        cgh.parallel_for(
+            sycl::nd_range<1>(grid_dims * block_dims, block_dims),
+            [=](sycl::nd_item<1> item_ct1) {
+                const int row = item_ct1.get_group(0);
+
+                top_k_radix_select_f32(
+                    src + (int64_t) row * ncols, dst_indices + (int64_t) row * k,
+                    (int) ncols, k,
+                    slm.get_multi_ptr<sycl::access::decorated::no>().get(),
+                    item_ct1);
+            });
+    });
+}
+
+// One work-group owns a whole row above, which leaves the device idle whenever a graph
+// has fewer rows than it has cores -- the common case at batch size 1, where the
+// sparse-attention indexer and the backend sampler both top-k a single row. The kernels
+// below spread one row over several groups instead.
+//
+// A digit pass now needs the whole row's histogram before any group can pick the pivot,
+// so the per-pass state moves to global memory and the passes become separate launches:
+// a work-group barrier no longer spans the row. Each group still accumulates into SLM
+// and contributes 256 global atomics at the end, so global traffic is per-group, not
+// per-element. The last group to finish a pass (the one whose fetch_add returns G - 1)
+// does the scan for the row and clears the histogram for the next pass, which keeps the
+// launch count at one per digit rather than two.
+//
+// Running all four digits unconditionally costs nothing in correctness: once a bucket
+// holds exactly the elements still needed, later digits only extend the prefix, and the
+// count of columns above that longer prefix grows by exactly as much as `need` shrinks.
+// The emit below therefore stays exact whatever pass the answer settled on.
+
+static constexpr int SYCL_TOP_K_RADIX_ROW_DONE   = SYCL_TOP_K_RADIX_BUCKETS + 0;
+static constexpr int SYCL_TOP_K_RADIX_ROW_PREFIX = SYCL_TOP_K_RADIX_BUCKETS + 1;
+static constexpr int SYCL_TOP_K_RADIX_ROW_MASK   = SYCL_TOP_K_RADIX_BUCKETS + 2;
+static constexpr int SYCL_TOP_K_RADIX_ROW_NEED   = SYCL_TOP_K_RADIX_BUCKETS + 3;
+static constexpr int SYCL_TOP_K_RADIX_ROW_CNT_GT = SYCL_TOP_K_RADIX_BUCKETS + 4;
+static constexpr int SYCL_TOP_K_RADIX_ROW_CNT_EQ = SYCL_TOP_K_RADIX_BUCKETS + 5;
+static constexpr int SYCL_TOP_K_RADIX_ROW_WORDS  = SYCL_TOP_K_RADIX_BUCKETS + 6;
+
+// How wide the split goes is a property of the device, not of the model: enough groups to
+// cover the cores, and no more. Past that the extra groups add histogram traffic without
+// adding bandwidth (measured on this device: 20 and 40 groups tie, 60 and 160 lose).
+//
+// nsm is max_compute_units / 16, i.e. it counts an Xe core as 16 EUs. That is a core's
+// width on Xe-HPG, but an Xe2 core is 8 XVEs wide, so on Battlemage the field reads half
+// the cores actually present (10 for a 20-core B60). The measured curve is flat from one
+// group per core to two and only falls off at three, so a factor of two covers the device
+// on Xe2 and lands in the flat region on Xe-HPG. It is the one number here that a correct
+// core count would remove; it was tuned on Xe2 and has not been measured on Xe-HPG.
+static constexpr int SYCL_TOP_K_RADIX_GROUPS_PER_NSM = 2;
+// Splitting trades one kernel for five. Below the width at which the single-group kernel
+// runs longer than those four extra launches, it wins on its own; measured break-even on
+// this device sits just under 64K columns.
+static constexpr int SYCL_TOP_K_RADIX_MIN_SPLIT_COLS = 65536;
+// A partition thinner than this cannot keep a group's sweep busy.
+static constexpr int SYCL_TOP_K_RADIX_MIN_PART_COLS  = 4096;
+
+static int top_k_radix_split_groups(const int device, const int64_t ncols, const int64_t nrows) {
+    const int64_t target = (int64_t) SYCL_TOP_K_RADIX_GROUPS_PER_NSM * ggml_sycl_info().devices[device].nsm;
+
+    // One group per row already, so a graph with rows enough to cover the device gains
+    // nothing from splitting and would only pay the extra launches.
+    if (ncols < SYCL_TOP_K_RADIX_MIN_SPLIT_COLS || nrows >= target) {
+        return 1;
+    }
+
+    const int64_t by_rows = target / nrows;   // floor: never overshoot a row that is nearly covered
+    const int64_t by_cols = ncols / SYCL_TOP_K_RADIX_MIN_PART_COLS;
+
+    return (int) std::max<int64_t>(1, std::min(by_rows, by_cols));
+}
+
+using top_k_radix_gatomic = sycl::atomic_ref<uint32_t, sycl::memory_order::relaxed,
+                                             sycl::memory_scope::device,
+                                             sycl::access::address_space::global_space>;
+
+static void top_k_radix_split_pass_f32(
+    const float *   src,
+    uint32_t *      state,
+    const int       ncols,
+    const int       k,
+    const int       shift,
+    const bool      first,
+    const int       part,
+    const int       nparts,
+    uint32_t *      slm,
+    const sycl::nd_item<1> & item_ct1
+) {
+    using local_atomic = sycl::atomic_ref<uint32_t, sycl::memory_order::relaxed,
+                                          sycl::memory_scope::work_group,
+                                          sycl::access::address_space::local_space>;
+
+    const int tid        = item_ct1.get_local_id(0);
+    const int block_size = item_ct1.get_local_range(0);
+
+    uint32_t * hist   = slm;
+    uint32_t * s_last = slm + SYCL_TOP_K_RADIX_HIST_SIZE;
+    uint32_t * s_row  = slm + SYCL_TOP_K_RADIX_HIST_SIZE + 1;   // prefix, mask, need
+
+    // The previous launch is the barrier that publishes these, so a plain load is enough.
+    // One lane reads them and the group takes them from SLM: a device-scope atomic load
+    // is uncached here, and having every work-item issue three of them off the same
+    // address costs more than the whole sweep below.
+    if (tid == 0) {
+        s_row[0] = first ? 0u : state[SYCL_TOP_K_RADIX_ROW_PREFIX];
+        s_row[1] = first ? 0u : state[SYCL_TOP_K_RADIX_ROW_MASK];
+        s_row[2] = first ? (uint32_t) k : state[SYCL_TOP_K_RADIX_ROW_NEED];
+    }
+
+    for (int i = tid; i < SYCL_TOP_K_RADIX_HIST_SIZE; i += block_size) {
+        hist[i] = 0;
+    }
+    item_ct1.barrier(sycl::access::fence_space::local_space);
+
+    const uint32_t prefix = s_row[0];
+    const uint32_t mask   = s_row[1];
+    const uint32_t need   = s_row[2];
+
+    const int copy  = tid & (SYCL_TOP_K_RADIX_HIST_COPIES - 1);
+    const int chunk = (ncols + nparts - 1) / nparts;
+    const int col0  = part * chunk;
+    const int col1  = std::min(ncols, col0 + chunk);
+
+    for (int col = col0 + tid; col < col1; col += block_size) {
+        const uint32_t key = top_k_radix_key(src[col]);
+        if ((key & mask) == prefix) {
+            const uint32_t bucket = (key >> shift) & (SYCL_TOP_K_RADIX_BUCKETS - 1);
+            local_atomic(hist[bucket * SYCL_TOP_K_RADIX_HIST_COPIES + copy]).fetch_add(1u);
+        }
+    }
+    item_ct1.barrier(sycl::access::fence_space::local_space);
+
+    // One global atomic per bucket per group, not per element.
+    for (int b = tid; b < SYCL_TOP_K_RADIX_BUCKETS; b += block_size) {
+        uint32_t sum = 0;
+        for (int c = 0; c < SYCL_TOP_K_RADIX_HIST_COPIES; c++) {
+            sum += hist[b * SYCL_TOP_K_RADIX_HIST_COPIES + c];
+        }
+        if (sum) {
+            top_k_radix_gatomic(state[b]).fetch_add(sum);
+        }
+    }
+
+    // Publish this group's bins, then claim the scan if this group is the row's last.
+    // The group-wide barrier flushes the atomics above; only the claiming lane needs the
+    // release, so the device-scope fence is paid once per group rather than per work-item.
+    item_ct1.barrier(sycl::access::fence_space::global_and_local);
+    if (tid == 0) {
+        sycl::atomic_fence(sycl::memory_order::release, sycl::memory_scope::device);
+        sycl::atomic_ref<uint32_t, sycl::memory_order::acq_rel, sycl::memory_scope::device,
+                         sycl::access::address_space::global_space> done(state[SYCL_TOP_K_RADIX_ROW_DONE]);
+        *s_last = (done.fetch_add(1u) == (uint32_t) (nparts - 1)) ? 1u : 0u;
+    }
+    item_ct1.barrier(sycl::access::fence_space::local_space);
+
+    if (*s_last == 0u) {
+        return;
+    }
+    sycl::atomic_fence(sycl::memory_order::acquire, sycl::memory_scope::device);
+
+    // Lane t takes bucket 255 - t, so an inclusive scan counts down from the top bucket.
+    uint32_t cnt = 0;
+    if (tid < SYCL_TOP_K_RADIX_BUCKETS) {
+        cnt = top_k_radix_gatomic(state[SYCL_TOP_K_RADIX_BUCKETS - 1 - tid]).load();
+    }
+    const uint32_t incl = sycl::inclusive_scan_over_group(item_ct1.get_group(), cnt, sycl::plus<uint32_t>());
+
+    if (tid < SYCL_TOP_K_RADIX_BUCKETS && incl >= need && incl - cnt < need) {
+        const uint32_t digit = (uint32_t) (SYCL_TOP_K_RADIX_BUCKETS - 1 - tid);
+        top_k_radix_gatomic(state[SYCL_TOP_K_RADIX_ROW_PREFIX]).store(prefix | (digit << shift));
+        top_k_radix_gatomic(state[SYCL_TOP_K_RADIX_ROW_MASK]).store(
+            mask | ((uint32_t) (SYCL_TOP_K_RADIX_BUCKETS - 1) << shift));
+        top_k_radix_gatomic(state[SYCL_TOP_K_RADIX_ROW_NEED]).store(need - (incl - cnt));
+    }
+    item_ct1.barrier(sycl::access::fence_space::local_space);
+
+    // Clear for the next pass; the next launch is the barrier that orders this.
+    for (int b = tid; b < SYCL_TOP_K_RADIX_BUCKETS; b += block_size) {
+        top_k_radix_gatomic(state[b]).store(0u);
+    }
+    if (tid == 0) {
+        top_k_radix_gatomic(state[SYCL_TOP_K_RADIX_ROW_DONE]).store(0u);
+    }
+}
+
+static void top_k_radix_split_emit_f32(
+    const float *   src,
+    int32_t *       dst_idx,
+    uint32_t *      state,
+    const int       ncols,
+    const int       k,
+    const int       part,
+    const int       nparts,
+    uint32_t *      slm,
+    const sycl::nd_item<1> & item_ct1
+) {
+    using local_atomic = sycl::atomic_ref<uint32_t, sycl::memory_order::relaxed,
+                                          sycl::memory_scope::work_group,
+                                          sycl::access::address_space::local_space>;
+
+    const int tid        = item_ct1.get_local_id(0);
+    const int block_size = item_ct1.get_local_range(0);
+
+    uint32_t * s_gt      = slm;
+    uint32_t * s_eq      = slm + 1;
+    uint32_t * s_base_gt = slm + 2;
+    uint32_t * s_base_eq = slm + 3;
+
+    uint32_t * s_row = slm + 4;   // prefix, mask, need
+
+    if (tid == 0) {
+        *s_gt = 0;
+        *s_eq = 0;
+        s_row[0] = state[SYCL_TOP_K_RADIX_ROW_PREFIX];
+        s_row[1] = state[SYCL_TOP_K_RADIX_ROW_MASK];
+        s_row[2] = state[SYCL_TOP_K_RADIX_ROW_NEED];
+    }
+    item_ct1.barrier(sycl::access::fence_space::local_space);
+
+    const uint32_t prefix = s_row[0];
+    const uint32_t mask   = s_row[1];
+    const uint32_t need   = s_row[2];
+
+    // Exactly k - need columns beat the pivot; the first need pivot-equal ones fill the tail.
+    const uint32_t base_eq = (uint32_t) k - need;
+
+    const int chunk = (ncols + nparts - 1) / nparts;
+    const int col0  = part * chunk;
+    const int col1  = std::min(ncols, col0 + chunk);
+
+    // Counting first and reserving one range per group keeps the row's two counters out of
+    // the inner loop: a per-element global atomic on a single address serialises the whole
+    // emit, and at k in the thousands that alone outweighs every read the kernel does.
+    for (int col = col0 + tid; col < col1; col += block_size) {
+        const uint32_t kp = top_k_radix_key(src[col]) & mask;
+        if (kp > prefix) {
+            local_atomic(*s_gt).fetch_add(1u);
+        } else if (kp == prefix) {
+            local_atomic(*s_eq).fetch_add(1u);
+        }
+    }
+    item_ct1.barrier(sycl::access::fence_space::local_space);
+
+    if (tid == 0) {
+        const uint32_t n_gt = *s_gt;
+        const uint32_t n_eq = *s_eq;
+        *s_base_gt = n_gt ? top_k_radix_gatomic(state[SYCL_TOP_K_RADIX_ROW_CNT_GT]).fetch_add(n_gt) : 0u;
+        *s_base_eq = n_eq ? top_k_radix_gatomic(state[SYCL_TOP_K_RADIX_ROW_CNT_EQ]).fetch_add(n_eq) : 0u;
+        *s_gt = 0;
+        *s_eq = 0;
+    }
+    item_ct1.barrier(sycl::access::fence_space::local_space);
+
+    const uint32_t base_gt_g = *s_base_gt;
+    const uint32_t base_eq_g = *s_base_eq;
+
+    for (int col = col0 + tid; col < col1; col += block_size) {
+        const uint32_t kp = top_k_radix_key(src[col]) & mask;
+        if (kp > prefix) {
+            dst_idx[base_gt_g + local_atomic(*s_gt).fetch_add(1u)] = col;
+        } else if (kp == prefix) {
+            const uint32_t pos = base_eq_g + local_atomic(*s_eq).fetch_add(1u);
+            if (pos < need) {
+                dst_idx[base_eq + pos] = col;
+            }
+        }
+    }
+}
+
+static void top_k_radix_split_f32_sycl(
+    ggml_backend_sycl_context & ctx,
+    const float * src,
+    int32_t * dst_indices,
+    const int64_t ncols,
+    const int64_t nrows,
+    const int k,
+    const int nparts,
+    dpct::queue_ptr main_stream
+) {
+    GGML_ASSERT(ncols <= INT32_MAX);
+    GGML_ASSERT(nparts > 1);
+
+    const int block_size = ggml_sycl_info().max_work_group_sizes[ctx.device];
+    GGML_ASSERT(block_size >= SYCL_TOP_K_RADIX_BUCKETS);
+
+    const size_t state_words = (size_t) nrows * SYCL_TOP_K_RADIX_ROW_WORDS;
+    ggml_sycl_pool_alloc<uint32_t> state_alloc(ctx.pool(), state_words);
+    uint32_t * state = state_alloc.get();
+
+    // Zero histogram, done counter and both emit counters. prefix/mask/need are seeded by
+    // the first pass, which ignores the stored values.
+    // The queue is in-order, so the passes below are already ordered after this fill.
+    SYCL_CHECK(CHECK_TRY_ERROR(main_stream->memset(state, 0, state_words * sizeof(uint32_t))));
+
+    const sycl::range<1> block_dims(block_size);
+    const sycl::range<1> grid_dims(nrows * nparts);
+
+    bool first = true;
+    for (int shift = 32 - SYCL_TOP_K_RADIX_BITS; shift >= 0; shift -= SYCL_TOP_K_RADIX_BITS) {
+        const bool is_first = first;
+        first = false;
+        main_stream->submit([&](sycl::handler &cgh) {
+            sycl::local_accessor<uint32_t, 1> slm(sycl::range<1>(SYCL_TOP_K_RADIX_HIST_SIZE + 4), cgh);
+
+            cgh.parallel_for(
+                sycl::nd_range<1>(grid_dims * block_dims, block_dims),
+                [=](sycl::nd_item<1> item_ct1) {
+                    const int g    = item_ct1.get_group(0);
+                    const int row  = g / nparts;
+                    const int part = g % nparts;
+
+                    top_k_radix_split_pass_f32(
+                        src + (int64_t) row * ncols,
+                        state + (int64_t) row * SYCL_TOP_K_RADIX_ROW_WORDS,
+                        (int) ncols, k, shift, is_first, part, nparts,
+                        slm.get_multi_ptr<sycl::access::decorated::no>().get(),
+                        item_ct1);
+                });
+        });
+    }
+
+    main_stream->submit([&](sycl::handler &cgh) {
+        sycl::local_accessor<uint32_t, 1> slm(sycl::range<1>(8), cgh);
+
+        cgh.parallel_for(
+            sycl::nd_range<1>(grid_dims * block_dims, block_dims),
+            [=](sycl::nd_item<1> item_ct1) {
+                const int g    = item_ct1.get_group(0);
+                const int row  = g / nparts;
+                const int part = g % nparts;
+
+                top_k_radix_split_emit_f32(
+                    src + (int64_t) row * ncols,
+                    dst_indices + (int64_t) row * k,
+                    state + (int64_t) row * SYCL_TOP_K_RADIX_ROW_WORDS,
+                    (int) ncols, k, part, nparts,
+                    slm.get_multi_ptr<sycl::access::decorated::no>().get(),
+                    item_ct1);
+            });
+    });
+}
+
+void ggml_sycl_top_k_radix(
+    ggml_backend_sycl_context & ctx,
+    const float *   src,
+    int32_t *       dst_indices,
+    const int64_t   ncols,
+    const int64_t   nrows,
+    const int       k,
+    dpct::queue_ptr main_stream
+) {
+    const int nparts = top_k_radix_split_groups(ctx.device, ncols, nrows);
+    if (nparts > 1) {
+        top_k_radix_split_f32_sycl(ctx, src, dst_indices, ncols, nrows, k, nparts, main_stream);
+    } else {
+        top_k_radix_f32_sycl(ctx, src, dst_indices, ncols, nrows, k, main_stream);
+    }
+}

ggml/src/ggml-sycl/topk-radix.hpp

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "common.hpp"
+
+// The legacy implementation uses SLM to implement sorting and top_k selection.
+// SLM is limited to 128KB on Xe, which limits how much can be sorted to k<32.
+// After a k=8, the radix selection becomes beneficial for most cases, because
+// scan-merge has (block + 1) * k pairs of (value, index). Given normal sorting of nlog(n),
+// radix-select becomes beneficial quite early. This sets it to 8 - however, the other parameters
+// (columns and rows) may also be a driving factor.
+// We select the legacy implementation for k below this constant because the overhead of radix select
+// exceeds the benefit for very small problems
+constexpr int SYCL_TOP_K_SCAN_MERGE_MAX_K = 8;
+
+// Top-k of every row of src, k indices per row into dst_indices, in no particular order.
+// Picks between the one-group-per-row and the split-row kernel from the shape and the device.
+void ggml_sycl_top_k_radix(
+    ggml_backend_sycl_context & ctx,
+    const float *   src,
+    int32_t *       dst_indices,
+    const int64_t   ncols,
+    const int64_t   nrows,
+    const int       k,
+    dpct::queue_ptr main_stream);

tests/test-backend-ops.cpp

@@ -11298,6 +11298,30 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
         }
     }
 
+    // qwen4exp sparse-attention indexer: nrows = n_tokens/n_stream, so tg gives nrows==1.
+    // Sweep nrows to expose how much of the device a single row leaves idle.
+    for (auto cols : {8192, 32768, 131072}) {
+        for (auto nrows : {1, 2, 4, 8, 16, 32}) {
+            test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, 2048));
+        }
+    }
+    // backend sampler: one row of the vocab (llama-sampler.cpp top_k)
+    for (auto k : {20, 40}) {
+        test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {151936, 1, 1, 1}, k));
+    }
+
+    // short rows, many of them: MoE routing and group selection. The opposite corner from
+    // the indexer, and the one where a work-group per row is the wasteful choice.
+    for (auto cols : {2, 16, 128, 1024}) {
+        for (auto nrows : {1024, 8192}) {
+            for (auto k : {1, 2, 8, 16, 32}) {
+                if (k <= cols) {
+                    test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, k));
+                }
+            }
+        }
+    }
+
     for (auto nrows : {1, 4, 8, 16}) {
         for (auto cols : {128, 1024, 4096, 8192, 16384, 32768, 65536, 131072, 200000, 2000000}) {
             test_cases.emplace_back(new test_cumsum(GGML_TYPE_F32, {cols, nrows, 1, 1}));