b10956

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 Website: https://llama.app Attestations: https://github.com/ggml-org/llama.cpp/attestations/47355610 macOS/iOS: macOS Apple Silicon (arm64) macOS Apple Silicon (arm64, KleidiAI enabled) DISABLED macOS Intel (x64) iOS XCFramework Linux: Ubuntu x64 (CPU) Ubuntu arm64 (CPU) Ubuntu s390x (CPU) Ubuntu x64 (Vulkan) Ubuntu arm64 (Vulkan) Ubuntu x64 (ROCm 10.0) Ubuntu x64 (OpenVINO) Ubuntu x64 (SYCL FP32) Ubuntu x64 (SYCL FP16) Android: Android arm64 (CPU) Windows: Windows x64 (CPU) Windows arm64 (CPU) Windows arm64 (OpenCL Adreno) Windows x64 (CUDA 12) - CUDA 12.4 DLLs Windows x64 (CUDA 13) - CUDA 13.3 DLLs Windows arm64 (CUDA 13) - CUDA 13.4 DLLs Windows x64 (Vulkan) Windows x64 (OpenVINO) Windows x64 (SYCL) Windows x64 (ROCm 10.0) openEuler: DISABLED openEuler x86 (310p) openEuler x86 (910b, ACL Graph) openEuler aarch64 (310p) openEuler aarch64 (910b, ACL Graph) UI: UI

Source: llama.cpp Releases | 2026-09-14

Related