RadiK: Scalable and Optimized GPU-Parallel Radix Top-K Selection
Yifei Li, Bole Zhou, Jiejing Zhang, Xuechao Wei, Yinghan Li, Yingda Chen
TL;DR
This work tackles scalable top-k selection on GPUs, where prior merge-based approaches struggle when k and input size are large due to on-chip memory limits. It introduces RadiK, a radix-based two-phase top-k algorithm (Radix Select and Filter) augmented by an optimization framework that maximizes memory bandwidth through hierarchical atomics and a flush-efficient write buffer, plus batch-specific techniques like task rescheduling and on-demand offset padding. A lightweight adaptive scaling technique is proposed to mitigate adversarial input distributions, significantly improving robustness. Comprehensive evaluations show up to 2.5x speedups for non-batch and 4.8x for batch queries, with substantial gains in robustness (up to 2.7x under adversarial conditions), demonstrating scalable, high-throughput top-k suitable for large vocabularies and large-scale vector queries.
Abstract
Top-k selection, which identifies the largest or smallest k elements from a data set, is a fundamental operation in data-intensive domains such as databases and deep learning, so its scalability and efficiency are critical for these high-performance systems. However, previous studies on its efficient GPU implementation are mostly merge-based and rely heavily on the fast but size-limited on-chip memory, thereby limiting the scalability with a restricted upper bound on k. This work introduces a scalable and optimized GPU-parallel radix top-k selection that supports significantly larger k values than existing methods without compromising efficiency, regardless of input length and batch size. Our method incorporates a novel optimization framework tailored for high memory bandwidth and resource utilization, achieving up to 2.5x speedup over the prior art for non-batch queries and up to 4.8x speedup for batch queries. In addition, we propose an adaptive scaling technique that strengthens the robustness, which further provides up to 2.7x speedup on highly adversarial input distributions.