RTop-K: Ultra-Fast Row-Wise Top-K Selection for Neural Network Acceleration on GPUs
Authors: Xi Xie
Yuebo Luo
Hongwu Peng
Caiwen Ding
Venue: ICLR 2025 Poster (2025)
Keywords: row-wise topk selection
GPU
CUDA
Primary area: infrastructure, software libraries, hardware, systems, etc.

Abstract: Abstract Top-k selection algorithms are fundamental in a wide range of applications, including high-performance computing, information retrieval, big data processing, and neural network model training. In this paper, we present RTop-K, a highly efficient parallel row-wise top-k selection algorithm specifically designed for GPUs. RTop-K leverages a binary search-based approach to optimize row-wise top-k selection, providing a scalable and accelerated solution.
We conduct a detailed analysis of early stopping in our algorithm, showing that it effectively maintains the testing accuracy of neural network models while substantially improving performance. Our GPU implementation of RTop-K demonstrates superior performance over state-of-the-art row-wise top-k GPU implementations, achieving an average speed-up of up to 11.49× with early stopping and 7.29× without early stopping. Moreover, RTop-K accelerates the overall training workflow of MaxK-GNNs, delivering speed-ups ranging from 11.97% to 33.29% across different models and datasets.

Published as a conference paper at ICLR 2025
RTOP-K: ULTRA-FAST FOR NEURAL NETWORK ACCELERATION ON GPUs
TOP-K SELECTION
'Xi Xie
'Yuebo Luo
University of Connecticut xi wwi.eeccoonnncommm edu
University of Minnesota Twin Cities 1uo004 66@umn edu
"Hongwu Peng
*Caiwen Ding
University of Connecticut
University of Minnesota Twin Cities di dingeeumn. edu
hongwa.peng@uunn www.:eeennne..cnn.ccmo
ABSTRACT
Top-k selection algorithms are fundamental a wide range of applications, in- cluding high-performance computing, information retrieval, big data processing, and neural network model training. In this paper, we present RTop-K, a highly ef- ficient parallel row-wise top-k selection algorithm specifically designed for GPUs. RTTop-K leverages a binary search-based approach to optimize row-wise top-k se- lection, providing scalable and accelerated solution. We conduct a detailed anal- ysis of carly stopping in our algorithm, showing that it effectively maintains the testing accuracy of neural network models while substantially improving perfor- mance. Our GPU implementation of RTop-K demonstrates superior performance over state-of-the-an row-wise top-k GPU implementations, achieving an average speed-up of up to 1 with early stopping and 7.29x without early stopping. Moreover, RTop-K accelerates the overall training workflow MaxK- de- livering speed-ups ranging from 11.97% to 33.29% across different models and datasets.
The GPU implementation can be found on Github
INTRODUCTION
Top-k selection is a classic algorithmic challenge that involves identifying the k largest or small- est elements from n input elements based on some predefined ranking criteria. The top-k selection algorithm has been widely applied in many traditional scenarios, such as high-performance comput- ing (HPC) (Muneer 2021), information retrieval (IR) (Ding Suel 2011). big data (Gaihre al., 2019), and data mining (Malkov & Yashunin 2018). Today, the top-k algorithm is increasingly ap- plied in the training and inference of neural network models. For example, the Avg-TopK (Ozdemir, 2023 pooling method has achieved more successful results in image classification accuracy and transfer learning models compared to traditional methods. TopK-SGD (Shi et al., 2019) applied to gradient sparsification techniques significantly reduces the communication traffic without obvious impact on the model accuracy. Combining top-k with sparse training (Jayakumar et al. 2021) can maintain constant sparsity and perform well while reducing resource requirements. In study (Cui et al. 2021). top-k attention loss function was introduced to address the top-k ranking prediction problem.
Graph Neural Networks (GNNs) have drawn tremendous attention in the past years due to their unique ability to extract latent information from graph data (Hu al.) 2020). In the design and acceleration of GNN training and inference, GPU platforms have become the prevalent choice due their multiple advantages. Firstly, compared to other processing hardware, GPUs provide superior processing power and memory throughput (Li et al. 2018). For example, the NVIDIA A6000 GPU boasts an impressive computation capability of 38.7 Tera FLOPS and memory throughput 768
"These authors contributed equally. + *Corresponding author.
https::/ / thub. com xi exi5) / opk

---

Published as a conference paper at ICLR 2025
GB/s. Secondly, many leading supercomputers (such Aurora and Eagle (LiveScience) 2023 use GPUs as their primary computing resource. Thirdly, many applications and services related to deep learning and neural networks developed and deployed on GPU platforms. However, GNN training and inference still typically pose strict challenges on latency and throughput (Xie et al, 2023
Recently, MaxK-GNN Peng et al. 2024 has achieved great success in the acceleration and opti- mization of GNN training and inference on GPUs. As shown in Figure [ this work introduces row- wise top-k selection before the feature aggregation step GNNs, which not only provides nonlinearity in GNNs to optimize the model's expressive ability but also demonstrates that performing SPMM oper- ations in GNNs with the row-wise right hand matrix can achieve several times speedup over traditional workflows while maintaining good model accuracy. The top-k selection operation in Max-GNN necessitates performing large-scale row-wise top-k computation, i.e., executing top-k operations simultaneously across batch of vectors on GPUs.
1. Top-k Selection
Raw-niis
Top-k
2. Feature aggregation Feature
Output
Figure 1: The core operation of MaxK- GNN, which introduces row-wise top-k se- lection into the GNN workflow to provide non-linearity and acceleration.
Traditional top-k algorithms and their GPU imple- mentations (Gaihre et al.) 202T, Zhang et al. 2023
Li et al. 2024) are typically optimized for single queries or limited batched queries, that is, for a large vector or small batch of large vectors (typically with a batch size not exceeding 100). However, the optimization focus for traditional scenarios differs from the row-wise top-k algorithm required for GNN training and inference. Implementing and optimizing row-wise top-k algorithms on GPUs pose challenges in terms dispersion, parallelism, and efficiency. Since row-wise top-k involves performing top-k operations on large batch of vectors simultancously, and each vector's length corresponds to the hidden dimension of the neural network layer (which usually does not exceed 1024), it is crucial to allocate only small and appropriate amount of GPU resources for each vector. Under these limited resource constraints, the various optimization methods proposed for large vectors in traditional top-k implementations may be overly complex and inefficient. We should seek simple and efficient algorithms tailored to this scenario.
Additionally, must consider the requirements and characteristics of applying row-wise top-k in neural networks. We only need to select the values of the top-k elements in each row and their indices in the vector. We do not need to perform sorting all; neither the selected elements in each row nor the remaining elements require sorting. Furthermore, given the neural network's s tolerant and robust nature, we can explore the feasibility of approximate top-k to further accelerate the overall algorithm.
To efficiently implement row-wise top-k on GPUs for neural network applications, we introduce RTop-K, highly efficient parallel top-k selection algorithm designed for large batch of limited- size vectors, with the capability of approximation to further enhance the speed of the row-wise top-k algorithm without compromising the accuracy of the neural network model.
We summarize our contributions as follows:
We provide a comprehensive summary GPU top-k selection algorithms and analyze the performance limitations of ofate-of-the-arr GPU implementations in the row-wise top-k selection scenario.
We propose binary search-basee top-k selection algorithm and provide theoretical anal- ysis of the effects of early stopping.
We implement the binary search-based top-k selection algorithm the GPU and con- duct comprehensive tests, demonstrating that it outperforms state-of-the-aa row-wise top-k GPU implementations, with early stopping having minimal impact testing accuracy.
2

---

SM). The usage of registers can affect the number Published as a conference paper at ICLR 2025
2 PRELIMINARY AND RELATED WORKS
2.1 TOP-K ALGORITIIMS
The heap- baap-based top-k algorithm (Cormen et al., 2009) uses a min-heap to maintain the top-k ele- ments by comparing and replacing the heap root when larger element encountered. QuickSelect (Dashti et al. 2013) leverages partition-based approach similar to quicksort to find the k-th largest element. The bucket-based algorithm (Yang et al. 2024) divides data into buckets, sorting only rel- evant ones to find the top-k elements, which is effective for uniformly distributed data. RadixSelect (Alabi et al., 2012) uses digit-wir wise sorting to identify the top-k elements efficiently. The bitonic top- k algorithm (Shanbhag et al. 2018) employs bitonic sorting to merge and find the top-k elements in parallel.
When considering these algorithms, must take into account their suitability for GPU implementa- tion and the optimization requirements for specific problem scenarios. For example, the heap-based top-k algorithm is not well-suited for parallelization structure operations and element-wise comparisons and swaps. Although QuickSelect, RadixSelect, and the bitonic top-k algorithm can be successfully implemented on GPUs, they still require con- siderable data access and resource demands when operating on vector. This makes it difficult to optimize for row-wise top-k scenarios, where a large batch of limited size vectors requires top-k selection simultaneously. necessitating simplified operations and limited resource usage per vector. The bucket-basee based top-k algorithm more friendly to row-wise top-k scenarios but still requires further simplification to enhance performance.
GPUs because it relies complex tree
2.2 GPU ARCHITECTURE
The architecture NVIDIA GPUs consists of an array of multithreaded Streaming Multiprocessors (SMs) designed to execute thousands of threads concurrently. A function that runs on a GPU is called a kernel.
Thread and Memory Hierarchy. NVIDIA GPUs organize threads into warps, with each warp containing 32 threads that execute the same instruction simultaneously. Warps are grouped into blocks, which reside the same Streaming Multiprocessor (SM) and can communicate via shared memory, a fast on-chip memory space. Blocks are further grouped into grids for specific kernel launches. Threads access data from multiple memory spaces: device memory (large but slower, accessible by all threads), shared memory (low-latency, for communication within a block), and registers (fastest, partitioned among threads of blocks that can be active on an SM.
Warp-Level Primitives. Warp-level primitives are a of operations that allow threads within warp to cooperate and communicate efficiently. These include:
Synchronization primitive: Ensures that all threads reach the same point in execution before proceeding.
Shuffle primitive: Allows threads to exchange values within a warp.
Ballot primitive: Enables threads to collectively determine which threads meet specified condition by generating a mask representing the threads that satisfy the condition.
Counting primitive: Counts the number of set bits in given mask, often used in conjunc- tion with the ballot primitive.
The flexible use of warp-level primitives is crucial for designing performrformnce kernels, as the efficiency of information sharing through these primitives can even surpass that of using shared memory.
2.3 GPU TOP-K INPLEMENTATIONS
Dr. Top-k et al. 2021 is delegate-centris system that helps reduce the workload of GPU top-k computations, including Radix Select, Bucket Select, and Bitonic Select. It achieves this by
3

---

Published as a conference paper at ICLR 2025
dividing the input into sub-ranges and selecting delegates from them, along with performing multi- GPU optimizations. A work et al.) 2023) proposed two optimization methods, AIR Top-k and GridSelect. AIR Top-k employs iteration-ussed design and adaptive strategy to minimize CPU-GPU communication and memory traffic, while GridSelect uses a shared queue and parallel two-step insertion decrease costly operations, enhancing parallel top-k efficiency on GPUs. A cent RadixSelect implementation RadiK (Li et al. (2024) utilizes an optimization framework tailored for high memory bandwidth and resource utilization, along with an adaptive scaling technique for enhanced robustness, that supports larger k values with high efficiency.
However, the above state-of the GPU implementations are optimized for limited batches of large vectors. For instance, Dr. Top-k, AIR Top-k, and RadiK are designed for scenarios where the vector size on the order of 22 (about one million elements), and the batch size does not exceed 100. This is not suitable for row-wise top-k applications, where the typical vector size is less than 1024, and the batch size can reach millions.
PyTorch's top-k implementation (Pytorch 2024) is suitable for wise top-k operations. It uses RadixSelect as the underlying method, which, as analyzed in Section 2.1, is overly complex for each limited-size vector. Moreover, its selection results are sorted, which is also unnecessary. Although it can handle large batch sizes, its efficiency is suboptimal in scenarios where a minimalistic top-k selection is critical for each single vector.
3 RTOP-K FRAMEWORK
Algorithm 1 Binary Search-based Top-k Se-
lection Algorithm
The row-wise top-k operation involves finding the largest (or smallest) & elements and their indices in each row of a matrix. Without loss of gener- ality, we assume finding the largest & elements. Suppose a matrix has N rows and M columns; the problem equivalent to performing top-k se- lection on N vectors of length M simultaneously. Since N can be extremely large and M is lim- ited, need to apply simplified algorithm to each vector, ensuring that the algorithm can exe- cute quickly with very limited computational sources and memory access. We adopt a binary search-based top-k algorithm, which even more convenient to execute than the bucket top-k algo- rithm.
Require: Vector v of size M. integer k
Ensure: Top-k largest elements and their in- dices in
1: min + min(v)
2: maa 4- max(v)
3: c+ étc.maa
4: ent +
5: while max min > e do
6: thres
7: ent + Hi 2 thres]]
8: if ent then
9:
max + thres
else if cnt > then
11: min t thres
12: else
13:
break
3.1 BINARY SEARCH- BASED TOP-K SELECTION ALGORITHM
14: end if
15: end while
16: if cnt < k then
The algorithm, as illustrated in Fig. 2 first re- trieves the min and max values of the vector, and then uses several iterations of binary search determine a threshold Uhres. The min, max, and Whres values are updated in each iteration, and the loop terminates when the number of ele- ments filtered by the current threshold equals k. One corner case that must be considered when v contains multiple equal or very close elements, and these elements happen to be near the "bor- derline" during top-k selection. In this scenario, relying solely on thres for selection makes it dif- ficult impossible to extract exactly i elements. It can be verified that in such cases, these borderline elements are always located between min and thres. Therefore, a second filtering step using min is sufficient supplement the selection and ensure that exactly A elements chosen.
17: elems,indices < {(v,,t) v; 2
thres)
18: Append the first k ent pairs of {(v., min 5 v. < thres} to
elems, indices
19: else
20:
elems, indices + the first pairs of [(v., i) v, 2 thres}
end if
22:
return elems, indices
4

---

Published as a conference paper at ICLR 2025
matrix W
fter
Min
Threshold
Max
nutm of mades
Iter
k
Iteration
Iter
ent k
iter 4
k
hidden dim
Figure 2: Illustration of the binary search-based top-k selection algorithm.
Table 1: Cumulative percentage of iterations where the loop exits for different & values (e 10 M = 256). Results are based on 10" repeated experiments for each k.
Iteration
k=16
k=32
k=64
k=96
k=128
3
4.1 4.13%
2.71%
1.96%
1.34%
1.58%
4
8.98%
5.32%
3.52%
3.00%
2.81 %
5
17.90%
10.64%
6.92%
5.84%
5.59%
6
33.86%
21.40%
13.87%
11.72%
11.15%
7
54.43%
38.84%
27.12%
23.29%
22.11%
8
72.38%
59.17%
46.64%
41.35%
39.93%
9
84.53%
76.00%
66.21%
61.48%
60.35%
10
91.88%
86.81%
80.68%
77.37%
76.62%
11
95.81% %
93.03%
89.79%
87.64%
87.18%
12
97.89%
96.45%
94.70%
93.57%
93.31% %
13
98.97%
98.21%
97.35%
96.70%
96.60%
14
99.52%
99.12%
98.67%
98.34%
98.25%
15
99.76%
99.53%
99.34%
99.20%
99.17%
16
100.00%
100.00%
100.00%
100.00%
100.00%
Average Exit
7.60
8.29
8.95
9.52
9.60
The parameter represents the precision of the algorithm, which determines the maximum width of the borderline and imposes an upper bound on the number iterations in the loop. When is set to O (which is effectively equivalent to the minimal distinguishable precision of the float type), Algorithm becomes an exact top-k selection algorithm.
Algorithm Binary Search-based Top-k Se- lection Algorithm with Early Stopping
Table presents the statistical results of the itera- tion counts at which the algorithm exits for differ- ent values of k with the vector's size M 256. For each & value, 10 repeated experiments were conducted, with the vector initialized with nor- mally distributed elements. It can be observed that the average iteration exit ranges from 7.6 to 9.6, and the probability of the iteration count being less than or equal 13 exceeds 95%.
Require: Vector v size M, integer k, inte-
ger maz åter
Ensure: Top-k largest elements and their in- dices in
1: min < min(v)
2: max +- max(v)
3: for iter +- to max iter do
4: thres maa
Algorithm I summarizes the complete binary basech-based top-k algorithm process. It contains number of branching conditions, and the loop length executed by each warp can be different. We attempt to further simplify it.
5: ent + iii
thres}]
6: if ent then
7: mar + thres
8: else
9: min thres
10: end if
Given the inherent robustness of neural networks, we explore the feasibility of incorporating early stopping into the algorithm. We present the pseudocode for the early stopping algorithm and then conduct a numerical analysis.
11: end for
12: clems, indices + ((v,, v;
min)
13: return first k pairs of elems, indices
As shown in Algorithm 2 the introduction of early stopping further simplifies the algorithm, with the main loop forcefully exiting in no more than mar iter iterations. The collection phase min the threshold, ensuring that only
5

---

Published as a conference paper at ICLR 2025
Table 2: Statistics of early stop top-k selection for Different & Values and Maximum Iterations (M=256). (1(%) represents the average relative error between the maximum element in early stop top-k selection and the maximum element in the optimal top-k selection. (2)% represents the average relative enror between the minimum element in early stop top-k selection and the minimum element in the optimal top-k selection. Hit( ()) represents the overlap ratio between the early stop top-k selection and the optimal top-k selection.
32
64
128
(
(i)))))
(S())
E
(1)(
(()
(x(
(5)
Eg(S)
llx(5)
()))
Eg(5)
[lx(%)
2
12.6
20.17
45.85
13.46
30.68
37.81
7.12
25.03
51.78
4.42
17.80
69.59
4.6
24.73
70.93
3
8.01
13.13
54.29
6.22
13.19
60.32
4.44
12.40
69.04
3.39
12.94
74.41
2.78
13.23
79.33
4
4.93
7.64
68.35
3,47
7.05
74.46
2.47
6.55
80.51
1.99
6.82
84.33
7.24
87.34
5
3,52
5.20
77.36
2.20
4,31
83.19
1.47
3,70
87.88
1.18
3.91
90.49
0.97
4.29
92.34
4.33
2.57
95.03
7
2.67
4.10
83.17
1.38
2.79
89.51
0.79
1,87
93.68
0.61
2.00
95.33
0.47
2.30
96.35
90.19
96.86
pass collection is needed. Table summarizes the hit rate (overlap ratio) and the average relative error between the early stopping top-k selection with different max_iter settings and the optimal top-k selection. The experiments were conducted with vectors of size M = 256 consisting of nor- mally distributed elements, and 10 repeated experiments for each condition. When maz_iter for h 2 32 (mar iter 2 6 for k = 16), both the maximum element and the minimum element in the early stopping top-k selection have an average relative error of no more than 5%. For 64, only iterations are needed for the hit rate between the early stopping top-k selection and the optimal top-k selection to exceed 80%. These results indicate that the early stopping top-k selection is numerically stable and controllable. We will further test the impact of early stopping top-k selection model accuracy in the experimental section.
Loading
Global
3.2 GPU IMPLEMENTATION DESIGN
Threads
Loop
Shared Memory
Both Algorithm and Algorithm are well- suited for GPU implementation, where sin- gle warp processes single vector of size M. Fig. 3 illustrates the GPU implementation de- sign, which can be divided into three stages: loading, searching, and selecting.
Searching
Max
Min
Mix
Min
Misc
Min
Mas
Min
Round
Loading stage: In this stage, each vector is loaded from global memory into the corre- sponding shared memory, which is done by warp looping through it. A synchronization barrier is set at the end of this stage.
Min
Masx
bitm
Round
Min
Broadcast
Searching stage: In this stage, each vector is handled by a single warp, assuming the warp contains threads (Fig. illustrates while in actual hardware environments w = 32). The key point of the implementation is that the mar, min, and !hres values at the begin- ning and in each iteration need be synchro- nized across all threads within the warp. This can be achieved using a combination of clas- sic tree-reduction and broadcast primitives. The first step is to obtain the mar and min of the vector. The vector is divided into [M/w] parts, with each thread responsible for extracting the maximum and minimum elements within its signed part. Then, a tree-reduction using the shuffle primitive is performed in log, steps
Coflect
4,
Cnt
Cnt
Cnt
Cntt
Reund
Reund2
Itee1
Selecting
Baller & Pop Cat
Shared
Global Nemory
0.8
0.7
0.5
0.9
Elemrats
Indlices
Figure 3: GPU implementation of the binary search-based top-k selection algorithm.

---

Published as a conference paper at ICLR 2025
max, Iter=2
max iterm4
max, Iter6
max, Iterws
Pytorch
Iter=3
Iter=?
stopping
N N=214, Mm256, 8.46x
N Mm512, 6.96x
N N=24 M-766, 5288x
Time (m84)
0.0
0.0
0.0
N= M=256, 921x
N=2 M=512, 7.88x
N=2 M=768, 6.15x
Time (md)
N-218, M-256, 8.92x
N-2 M-512, 7.11x
N-2 M-768, 5,71x
Tima (m4)
K
N= M=256, 8.95x
N= M=512, 7.15x
N= M=768. 5.67X
Time [md]
50
1.6
32 64
64
96
128
Figure 4: Comparison of kernel execution time (ms) between RTop-K with different early stopping maa jier values and without early stopping (c 10-16), against PyTorch various configurations
of (N, M, k), where N 214, 216 218 2nd .M 256, 512, 768, and 16, 32, 64, 96, 128. The average speedup of the no early stopping version for each (N, M) setting is indicated in the title of each subplot.
to obtain the maximum and minimum elements across the entire warp, and these values are broad- casted to all threads within the warp. The second step is to perform binary search to determine the selection threshold. each iteration, the count of elements exceeding the current threshold is accumulated and broadcasted using the to simultaneously update the mar, min, and thres values. Once the exit condition is met or the maximum number of iterations is reached, the final threshold is obtained.
tree-reduction method. Then, each thread uses end
Selecting Stage: A single warp performs a one-pass or two-pass selection over each vector. Let ent be the number of elements greater than equal to the final threshold thres. The selection applies two conditions separately: selecting elements where r 2 thres, and, if needed, selecting additional elements where thres > x 2 min.
If cnt 2 k. the first condition is sufficient to produce & elements, and the second condition is skipped. Otherwise, elements satisfying the first condition are selected, and the remaining k k-ent elements are supplemented using the second condition.
To implement this efficiently, we use the ballot primitive merge and broadcast the selection masks across threads within a warp. The popent primitive is then used to compute the inclusive prefix sum of selected elements for each thread. Based on the prefix sum, each thread finalizes its selection decision: if the prefix sum exceeds kt its selection is disabled. The values and indices of the final selected elements are written to the output buffer in global memory.
This design requires data writes outside of registers, except for loading the vector and dumping the results. During the searching and selecting stages, warp-level primitives are utilized to achieve highly optimized inter-thread collaboration. Moreover, each warp operates independently in parallel, maintaining high overall efficiency.
EXPERIMENTS
4.1 SETUP AND CONFIGURATION
The CUDA source code of RTop-K is compiled using NVCC, version 12.6, and executed NVIDIA A6000 platform running Ubuntu 22.04. We conduct comprehensive performance tests on the RTop-K kernel, covering various input matrix dimensions, with the number of rows N ranging

---

Published as a conference paper at ICLR 2025
Table 3: Average speed up of RTop-K compared to PyTorch implementation (c = 10-16 1G for No
Early Stopping) across different M values.
Max Iteration
2
3
4
5
6
7
8
No Early Stopping
M=256
13.07
12.32
11.46
10.86
10.32
9.88
9.55
8.88
M=512
11.66
11.37
10.43
9.51
8.87
8.34
7.98
7.27
M-768
9.73
9.44
8.72
7.75
7.16
6.78
6.46
5.72
Average
11.49
11.04
10.20
9.37
8.79
8.34
7.99
7.29
Table 4: Graph data and the baseline testing accuracy of the MaxK-GNN based GNN model along with the percentage of time spent on row-wise top-k operations during training.
GNN Model
GraphSAGE
GCN
GIN
Graph
#Nodes
Acc(%)
Top-k Prop(%) (%) Acc(%)
Top-k Prop( %) Acc(%)
Top-k Prop(%)
Ogbn-products
2449029
80.08
19.81
76.6
19.61
77.77
19.67
Yelp 716847
61.09
26.09
48.26
25.84
43.16
25.92
Reddit 232965
96.74
11.66
95.18
11.61
94.96
11.62
Flickr 89250
53.44
26.86
50.42
26.78
51.73
26.73
from 214 to 2 hidden dimensions M ranging from 256 to 768, and k values ranging from 16 to 128. all cases, we evaluate the speed of RTop-K with different early stopping settings, including mar iter values from 2 to 8, as well pared against the row-wise top-k implementation provided by PyTorch (Pytorch] (002)), which is the state-of-the-ard row-wise top-k implementation supporting large number of The latency measurements are conducted using the Nsight Compute tool NVIDIA 2023).
early stopping (c 10 10-16). The results com-
We also integrate the RTop-K kernel into MaxK-GNN models to evaluate the overall speedup of the entire training process and the impact of different early stopping settings on test accuracy. The evaluation covers three models in MaxK-GNN: GraphSAGE (Hamilton et al. 2017), GCN (Kipf
& Welling 2016). and GIN (Xu et al. 2019). The graph datasets used include Flickr (McAuley & Leskovec 2012). Yelp (Zeng et al. 2020). Reddit (William L. Hamilton 2017). and Ogbn-products
(Hu et al.
2020
4.2
RToP-K KERNEL EVALUATION
Fig. presents a comprehensive time profiling of RTop-K compared to the PyTorch implementa- tion. It can be observed that RTop-K demonstrates significant speed improvements over PyTorch across various early stopping mar_iter settings. Even with no early stopping, RTop-K significantly outperforms Py PyTorch all scenarios. Moreover, the dimension that has the most impact on the speed-up ratio is M, while N and A have relatively smaller effects. Table summarizze the average speed-up RTop-K relative PyTorch for different values of M. When M = 256, the speed-up varies from 9.55x to 13.07x x across different mar matriter settings, and even with early stopping, the speed-up is still as high as 8. 8.88x. On average, the speed-up ranges from 7.99x to 11.49x with early stopping, and 7.29x with no early stopping.
We observe that the speed-up with no early stopping is close to that of max_iter = 8, which indicates that although the binary search requires many iterations in the worst-case, it typically exits early in most cases. This observation consistent with the results presented Table Even with few bad cases, the overall kernel speed remains unaffected.
4.3 MODEL TRAINING AND TESTING PERFORMANCE EVALUATION
Table summarizes the proportion of time spent row-wise top-k operations in several MaxK- GNN training instances. is evident that row-wise top-k operations account for a substantial portion of the training time, ranging from 11.61% in the Reddit GCN training to 26.86% in the Flickr GraphSAGE training. This indicates that optimizing row-wise top-k operations is meaningful for improving their training efficiency.
8

---

Published as a conference paper at ICLR 2025
Accuracy
Optimal Top KAccuracy
Overal Speed up Ratio
SAGE - Redditt
SAGE Flickr
SAGE Ogbn-products
SAGE Yelp
Aetuney
1.02
1.00
- 4
Speeed up Ratio
Aecuracy
0.56-
up Ratio
Sp-odd up Ratio
Accurecy
0.66
1.32
1.31
Speedd up- Rotto
1.17
1.21
0.54
0.80
1.20
1.09
0.52
1.33
1.29
3
Max
Iter
GCN Reddlt
GCN Flickr
GCM Oggn products
Yelp
1.02
8282
1.23
1.32
Accurcy
0.58
Spreed up Rasio
Accursoy
1.3 12
Spreesh up Ratio
Spaiddup up Ratio
Accurssy
0.50
1.35
Speeddon Relio
0.96
0.48
0.94
0.76
1.29
0.46
Max fter
Max Iter
Max iter
Max iter
GIN Reddit
GIN Flickr
GIN Ogbn-products
- Yelp
0.81
Arcuraty
S- de- do
1.12
Speed-in up Rabo
Accuracy
0.53
Speedt -pp Ratio
Accurncy
Spoed-uip Raalo
accuracy
1.31
Speed-up Ratoo
0.00
1.21
0.96
L10
0.52
0.53
0.78
1.20
0.44
0.94
0.50
1.19
1.20
0.42
1.2
Figure 5: Overall training speed-up ratio and testing accuracy of applying RTop-K to various MaxK- GNN model training processes on different graphs. Setting: N #Nodes, M 256, h: 32.
The impact of applying RTop-K on speed and accuracy with different early stopping settings in the actual training of these models is shown in Fig. 5 with the setting of N #Nodes, M = 256, and & 32. For all GNN models and all graphs, the application the overall training workflow. Specifically, under different max iter settings, the average overall training speed-up Reddit ranges from 11.97% % 12.21%. %2 for Flickr from 32.48% to 33.29%, for Ogbn-products from 22.00% to 22.74%, and for Yelp from 31.21% to 32.42%.
RTop-K K effectively accelerates
can be observed that the testing accuracy of the models remains high. Except for the GIN training on Flickr, the testing accuracy across different max iler settings for other cases fluctuates around the testing accuracy achieved with the optimal row-wise top-k selection. In many cases, applying early stopping for row-wise top-k selection even results in higher testing accuracy. This is also manifestation of the inherent robustness of GNNs.
CONCLUSION
In this paper, presented RTop-K, highly efficient parallel row-wise top-k selection algorithm for GPUs. By employing a binary search-based approach, RTop-K significantly accelerates top-k operations while maintaining the accuracy of neural network models, as confirmed by our theoretical analysis. Comprehensive kernel evaluations demonstrate that RTop-K K outperforms state-of-the-an GPÚ implementations, achieving an average speed-up of up to 11.49x with early stopping and 7.29x without early stopping. The evaluation of the overall MaxK-GN -NN training workflow with RTop-K shows that RTop-K provides overall speed-ups ranging from 11.97% to 33.29% across different models and datasets, with early stopping having almost no impact on testing accuracy.
6
ACKNOWLEDGMENTS
This research was supported in part by NSF SHF-2505770, Semiconductor Research Corporation (SRC) Artificial Intelligence Hardware program.

---

Published as a conference paper at ICLR 2025
REFERENCES
Tolu Alabi, Jeffrey D Blanchard, Bradley Gordon, and Russel Steinbach. Fast k-Selection Algo- rithms for Graphies Processing Units. Journal of Experimental Algorithmics, 17:4-2, 2012.
Thomas H Cormen, Charles E Leiserson, Ronald L. Rivest, and Clifford Stein. Introduction to Algorithms. MIT press, 2009.
Chaoran Cui, Jian Zong, Yuling Ma, Xinhua Wang, Lei Guo, Meng Chen, and Yilong Yin. Tri- branch convolutional neural networks for top-k focused academic performance prediction, 2021.
URL httpt:tttattxttt org/ /bs/2107.... 10424
Ali Dashti, Igor Komarov, and Raissa M D Souza. Efficient computation of k-nearest neighbour graphs for large high- -ighddimensiona data sets on gpu clusters. PLoS One, 8(9))e74113, 2013. doi: 10.1371 /11371/journll.pog pone.004111...
Shuai Ding and Torsten Suel. Faster Top-k Document Retrieval Using Block-Max Indexes. In Proceedings of the 34th International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 993-1002, 2011.
Anil Gaihre, Santosh Pandey, and Hang Liu. Deanonymizing Cryptocurrency with Graph Learning: The Promises and Challenges. In Conference Communications and Network Security (CNS), pp. 1-3. IEEE, 2019.
Anil Gaihre, Da Zheng, Scott Weitze, Lingda Li, Shuaiwen Leon Song, Caiwen Ding, Xiaoye S. Li, and Hang Liu. Dr. top-k: delegate-centric top-k on gpus. In Proceedings of the International Conference for High Performance Computing, Nefworking, Storage and Analysis, SC '21, New York. NY. USA, 2021. Association for Computing Machinery. ISBN 4781450384, doi: 10. (145/3458817.3476//
COLL:cts::/// :/dci. org 10 11w111153335 111111111/1111.11 3476141
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. Advances in neural information processing systems, 30, 2017.
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems, 33:221 11-22183,3 2020.
Siddhant M. Jayakumar, Razvan Pascanu, Jack W. Rae, Simon Osindero, and Erich Elsen. Top-kast: 23151
Top-k always sparse training, 2021. URL httpt::ttaaaiiiiiee /tt.a//tts////aaa....r////a/a/aa
Thomas N Kipf and Max Welling. Sumi-super classification with graph convolutional net- 2016.
works. arXiv preprint
Ang Li, Weifeng Liu, Linnan Wang. Kevin Barker, and Shuaiwen Leon Song. Warp-consolidation::
A Novel Execution Model for GPUs. In International Conference on Supercomputing, 2018.
Yifei Li, Bole Zhou, Jiejing Zhang, Xuechao Wei, Yinghan Li, and Yingda Chen. Radik: Scalable and optimized gpu-parallel radix top-k selection. In Proceedingy of the 38th ACM International Conference on Supercomputing, ICS '24, pp. 537-548, New York, NY, USA, 2024. Associa- tion for Computing Machinery. ISBN 9798400706103.
doi: 10.1145/365000033555699 URL
https://doi htt///di.org/0/00 wwt.or/.............11111 000...11//355000333333 365655999
LiveScience. wwwsttpwwwrfll----ooo hetsss/wwwwrrrllsuuuooooorrtttre
The most at
powerful supercomputers in the world right ence. com/technol loy//
now.
Available
https::/wwxz
2023. Accessed: 2024,
Sep 30.
Yu A Malkov and Dmitry Yashunin. Efficient and Robust Approximate Nearest Neighbor Search Using Hierarchical Navigable Small World Graphs. IEEE Transactions on Pattern Analysis and Machine Intelligence, 42(4))824-836, 2018.
Julian McAuley and Jure Leskovec. Image labeling on networkk Using netial- metadata for 59.
image classification. 07 2012. ISBN 978-3-642-33764422 doi: 10.
10

---

Published as a conference paper at ICLR 2025
Muneer. arrayfirerequest.
hettp::/ggrops.ccoos arzayiiee-ussss//ooooooTTaaaa 2021. Accessed: 2021, 17.
Available
NVIDIA. Nvidia nsight
2023. Accessed: 2023-08-20.
Hongwu Peng, Xi Xie, Kaustubh Shivdikar, Md Amit Hasan, Jiahui Zhao, Shaoyi Huang, Omer Khan, David Kaeli, and Caiwen Ding. Maxk-gnn: Extremely fast gpu kernel design for acceler- ating graph neural networks training. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, pp. 683-698, 2024.
Pytorch. topc
https::/pytore: topk htmll 2024. Accessed: 2024-08-15.
Anil Shanbhag, Holger Pirk, and Samuel Madden. Efficient Top-K Query Processing on Massively Parallel Hardware. In Proceedings of the 2018 International Conference on Management of Data, pp. 1557-1570. ACM, 2018.
Shaohuai Shi, Xiaowen Chu, Ka Chun Cheung, and Simon See. in distributed deep learning, 2019, URL
Understanding top-k sparsification 08772
Jure Leskovec William Hamilton, Rex Ying. Inductive representation learning on large graphs. arXiv preprint arXiv: 777702226666 2017.
Xi Xie, Hongwu Peng, Amit Hasan, Shaoyi Huang, Jiahui Zhao, Haowen Fang. Wei Zhang, Tong Geng, Omer Khan, and Caiwen Ding. Accel-gen: High-performance gpu accelerator design for graph convolution networks. In 2023 IEEE/ACM International Conference Computer Aided Design (ICCAD), pp. 01-09, 2023. doi: 10. N01109//CCADS5799222....000
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanic Jegelka. How powerful are graph neural
networks? In International Conference Learming Representations, 2019. URL https:
//openreview. /ttt//frummm wtw.:rrrmmmmmmyreeaeree
d=ryGs 611.G666.3 61 A5Km
Yiqing Yang, Guoyin Zhang, Yanxia Wu, Zhixiang Zhao, and Yan Fu. Split- partition (sbp): novel execution model for top-k and selection algorithms gpus. J. Supercomput, 80(11 15122-15160, March 2024. ISSN 0920-8542,
doi: 10.1007/s1 1220000/11127-0242222 1227-2240111-- URL.https.
7/doi. org/1.. 1007 000000//11121.0 024-06031-
Hanqing Zeng, Hongkuan Zhou, Ajitesh Srivastava, Rajgopal Kannan, and Viktor Prasanna. Graph- SAINT: Graph sampling based inductive learning method. In International Conference on Learn- ing Representations, 2020. URL (nttssts::/ http::::///////////eeeee net/ wtwp:evwwwwweeeeeeuenneeeeeeee.e hott://ErrummmdddBeeeepii
Jingrong Zhang, Akira Naruse, Xipeng Li, and Yong Wang. Parallel top-k algorithms on gpu: A
comprehensive study and new methods. In Proceedings of the International Conference for High
Performance Computing, Networking, Storage and Analysis, SC '23, New York, NY, USA, 2023. Association for Computing Machinery. ISBN 9798400701092.
doi: 10.1 00.114/335818843607666
URL https: / /10 1ee..1111111..1
3607062
Cüneyt Özdemir. Avg-topk: A new pooling method for convolutional neural networks. Expert Systems with Applications, 223: 19892, 2023. ISSN 0957-4174. doi: https://boi.rr//0.100666
j.eswa.023. URL https: www. www.wwwwiiinnciie........c...... com/ com/sciince/artioo.
041/8095744742300333
A THE EXPECTATION OF THE ITERATION COUNTS OF ALGORITHM I
Assume that a vector of length M has elements that follow a normal distribution N(u,a2) The probability that an element æ exceeds a threshold thres is given by P(x > thres)
1 - 4b [hres-p where & is the cumulative distribution function of the normal distribution. If the expected number of elements selected from v is k, the expectation of thres, denoted as E(thres), satisfies:
11

---

Published as a conference paper at ICLR 2025
M (1- ( )) =k
E(thres) μ x-1 (1
k
M
Considering the distinguishable interval ó between the k-th and (A: 1)-th largest elements, the length of this interval is:
M [(E(thres))
where
(E(thres) 11)2
f(E(thres))
exp
0 0 2m
20²
the probability density E(thres). The length of the initial search interval D is given by:
D mvx(v) min(v) = 20 2n
Each iteration of binary search halves the search interval length and moves closer to E(thres). The expected number of iterations En)) required for the algorithm to exit is determined by the search interval shrinking to within 8. Thus, E(n) can be approximated as:
D
E(n) loga
loga
20 2ev2lnMM M M f(E(thres))
8
In M
$ (1
h
2
lloa2
2M
(4)
2 l
M
We compared the calculation results of Equation with more detailed experimental results, as shown in T can be observed that the results match well, but F(n) always slightly larger than the measured average exit. This could be because the estimation of the initial search interval, D Rs 2 M, is valid only when M is sufficiently large. When M is not large enough, the lack of tail samples causes the actual initial search interval to be smaller.
B A COMPREHENSIVE ANALYSIS OF THE PERFORMANCE OF RTOP-K WHEN APPLIED VARYING VECTOR SIZES
Our design fixes warp to process vector. As the vector size M increases, the shared memory required per warp also increases. Given that the available shared memory per block has a limit, on the A6000 GPU, we allocate only [ warps per block. For M > 8192, our current shared memory-based acceleration strategy cannot be directly applied.
For M < 8192, the speedup of RTop-K K relative PyTorch shown in Figure
Considering the lower-bound speed of RTop-K (no early stopping version):
When M is below 1280, RTop-K achieves a 4.9x to 12.5x speedup over PyTorch.
When M is between 1280 and 3072, RTop-K achieves 2.3x to 4.9x speedup over Py- Torch.
When M is between 3072 and 6144, KTop-K achieves 1.1x to 2.3x speedup over Py- Torch.
When M is between 6144 and 8192, RTop-K is slower than Torch, with only the early stopping version using small mar_iter still being faster than PyTorch.
Theoretically, as shown in Equation the expected number of search iterations for Algorithm
is:
12

---

Published as a conference paper at ICLR 2025
Table Cumulative percentage iterations where the loop exits in Algorithm 1 for different M, values, with 0. couple, and theoretical values (n) are also provided. Iters
Experimental results are based 10t repeated experiments for each M,
M,
256
256
1024
1024
1024
1024
4096
4096
4096
4096
8192
8192
8192
8192
64
128
64
128
256
512
64
128
256
512
64
128
256
512
I
0.12
1.5
0
0
0.02
0.53
0
0
0
0
0
0
2
0.17
1.5
1.25
0.17
0.02
0.53
0.41
0.41
0.18
0.12
0.23
0.24
0.02
3
1.98
1.62
1.27
0.5
0.62
0.54
0.43
0.41
0.2
0.08
0.3
0.24
0.25
0.02
4
3,48
2.9
1.93
1.31
0.99
0.7
1.54
0.63
0.36
0.33
1.11
0.7
0.32
0.15
5
6.88
6.06
4.05
2.42
1.67
1.25
2.83
1.34
0.76
0.61
2.43
1.3
0.7
0.31
6
13.96
11.64
8.18
4.49
3.24
2.64
5.4
2.84
1.64
1.06
4.68
2.62
1.46
0.75
27.3
23.03
16.09
9.42
6.41
5.16
11.13
5.99
3.37
2.15
9.18
5.28
2.98
1.6
8
46.47
40.18 30.63
18.77
12.7
10.19
21.56
12.05
6.89
4.15
19.25
10.62
5.59
3.27
9
66.34
60.57 50.71 35.21
24.91
19.69
39.04
24.37
13.69
8.23
34.8
21.5
11.49
6.72
10
81.07
76.76 69.08 55.43 43.92
35.98
60.23 42.59
26.65
16.71
55.41
38.13
22.42
13.3
11
90.11 87.67 82.99 72.93 63.57
55.7 76.59 62.55
46.7
31.29
73.09
58.45
39.96 25.94
12
95.17 93.48 90.71 84.79 78.45 73.32 87.42 78.11 66.36 52.28 85.21 75.11 60.66 44.97
13
97.46 96.8 95.19 91.9 88.37 85.06 93.42 87.9 81.06 70.87 92.16 86.37 77.12 64.65
14
98.68 98.35 97.59 95.86 93.82 92.47 96.48 93.18 89.91 83.91 96.29 92.63 87.55 80.25
15
99.46 99.21 98.79 97.95 96.87 96.32 98.24 96.3 94.73 91.34 98.15 96.39 93.48 89.3
16
99.73 99.56 99.43 98.89 98.46 98.06 99.21 98.14 97.35 95.39 99.11 98.19 96.73 94.46
17
99.85 99.72 99.71 99.51 99.18 99.03 99.59 99.02 98.51 97.65 99.48 99.18 98.35 97.29
18
99.96 99.85 99.9 99.73 99.62 99.44 99.77 99.45 99.2 98.8 99.77 99.59 99.21 98.66
19
99.99 99.93 99.95 99.84 99.84 99.67 99.87 99.75 99.64 99.38 99.89 99.76 99.54 99.33
20
99.99 99.98 99.97 99.9 99.93 99.85 99.95 99.87 99.84 99.76 99.98 99.89 99.8 99.69
21
99.99 99.99 99.98 99.95 99.97 99.93 99.97 99.94 99.95 99.95 99.99 99.98 99.94 99.84
22
99.99 99.99 99.99 99.96 99.99 99.95 99.99 99.98 99.98 99.97 99.99 99.98
99.96 99.9
23
99.99
99.99 99.99 99.99 99.99 99.96 100
99.98 100
99.98 100
99.99 99.94 100
24
100
99.99 100
100
99.99 99.97
100
99.99
100 99.96
25
99.99
100 99.99
99.99
99.98
26
100
100
99.99
99.99
27
100
99.99
28
100
Avg
8.72
9
9.53
10.31
10.87
11.24
10.07
10.95
11.73
12.46
10.3
11.14
12.02
12.8
E(n)
9.08
9.41
9.87
10.62
11.24
11.57
10.36
11.2
12
12.75
10.54
11.41
12.26
13.06
13

---

Published as a conference paper at ICLR 2025
Speedup by and Selected max iter
20
iter
max
max,
No Early Stopping
10
Speedup
8
768
19
2048
3524
3968
4096
4224
435
44
360
486
299?
5128
24
500
3888
6016
6
078
6912
7040
7668
78
1638
7808
8066
8192
Figure 6: Speedup of RTop-K relative Py PyTorch for different vector sizes M and number of vectors
N 65536. The average speedup for each case is computed over k = 64, 128, 256, 512 and k < M.
Precision 10 is used for the no early stopping version.
In M
A M ))
2M V In M π
((n) log2
2M
Φ-1
log2
Olog M)
2ln
In each iteration, reduction operation of length M is required. Since warp is fixed for one vector, the time complexity of this reduction is equivalent to serial reduction, which is O(M). Therefore, the total time complexity of Algorithm [ is O(M log M). In contrast, PyTorch's un- derlying operation, RadixSelect, has a time complexity of O(M). Thus, when M is sufficiently large, Algorithm will lag behind traditional algorithms.
However, from a practical perspective, for M < 8192, as shown in Table 5, the growth of E(n) very slow. Additionally, since the search range has lower bound (depending the data type), the number of search iterations has an upper bound. Moreover, the searching stage is fully executed in shared memory, which leads to decreasing proportion of time spent the searching stage compared to the loading and selecting stages, as these involve increasing global memory accesses with the growth of M.
Therefore, we believe the actual time complexity of Algorithm I is less than O(M log M). As M increases, the relative speedup of RTop-K compared to PyTorch decreases primarily because relative efficiency of PyTorch's is RadixSelect improves (the proportion of its initialization, histogram construction, and indexing overhead decreases).
We also evaluated the performance of Algorithm with different precision settings, and the results are shown Figure
We found that precision has almost no impact speed. Even with the setting of 0, as shown in Table 5
Algorithm exits within iterations in the vast majority of cases, and the remaining rare cases have negligible impact the overall performance. This is also because the searching stage is fully executed in shared memory, making Algorithm [ less sensitive to the number of search iterations.
14

---

Published as a conference paper at ICLR 2025
Speedup by and Selected Precision, No Early Stopping
20
0.0001
10
Speedup
Ni00
140
153
1662
192
1922
2048
638
3583
3988
40
4227
43
44 2480
47
4864
4992
770
5888
2012
84
6917
7040
7168
555
6600
808
36
8064
8192
Figure 7: Speedup of RTop-K (no early stopping version) relative PyTorch for different vector
sizes M and different precisions, with the number of vectors N 65536. The average speedup for each case computed over 64, 128, 256, 512 and M.
15


Reviews:
[{"review_id": "vNfFcjF3hE", "reviewer": "Reviewer_3azA", "summary": "This paper presents RTop-K, a GPU acceleration of row-wise TopK selection. The proposed algorithm is based on binary search. Also, it introduces the early stopping strategy to improve the performance further. The experimental results show that it outperforms the state of the art row-wise TopK implementation.", "questions": "Please refer to Weakness above.", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well organized and easy to follow.\n2. This work addresses one of the key component in the GNN training.", "weaknesses": "1. Regarding the kernel performance comparison, the authors compares the solution with the PyTorch. The proposed solution should be compared with other topK solutions, like Shanbhag et al. [1], or radix select-based etc--assessment that \"not suitable for GNN is not sufficient\", in my opinion. Besides, I am not sure the recent PyTorch is the best baseline to compare, as I see some complaints that it is slower than torch.sort (I am not sure whether this is already resolved or not).\n\n[1] Anil Shanbhag, Holger Pirk, and Samuel Madden. Efficient Top-K Query Processing on Massively Parallel Hardware.\n\n2. Limited applicability: the experiment shows the speedup for Max K-GNN (includes GraphSAGE, GIN, GCN) in the only. Not sure if there's a broader impact on other DL networks.\n3. Theoretical analysis on early stopping.\n4. It'd be great if you can mark speedup on Fig. 4\n5. How the implementation is close from the roofline? Can you please share the Nsight profiling results?\n6. Is TopK bottleneck in the GNN? Can you please share the breakdown results, not just showing the portion of the TopK part?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6sPnbljYyU", "reviewer": "Reviewer_4z6F", "summary": "This paper looks at performing top-k row-wise selection in matrices that are, broadly, tall and skinny – with fewer than 1000 elements per row but perhaps millions of rows. Most existing algorithms focus on the opposite regime – relatively few rows that are much longer – and this distinction is important for designing performant GPU implementations. The authors introduce an algorithm that seeks to outperform alternative approaches, including the default RadixSelect algorithm in PyTorch, and is designed for the hardware, using warp-level parallelism within a row and careful use of the cache hierarchy. They also introduce an early stopping version of the algorithm that uses a fixed number of iterations, as load imbalance across warps/rows could significantly decrease the potential speedup. Results are shown in two settings. First, they consider the performance of the algorithm with and without early stopping on synthetic data, evaluating the number of iterations required and finding acceptable error rates when using early stopping. When compared with PyTorch’s implementation, they find substantial speedups (2-6x). Second, the implement their top-k in a few max-k GNN algorithms, like SAGE, where top-k calculations account for a non-trivial amount of runtime, and test on a few benchmark graph datasets. They report speedups of 10-30% and, with early stopping, find that model test accuracy fluctuates around that expect from using the optimal top-k values.", "questions": "The authors did a nice job of taking an HPC paper and providing a compelling presentation for an ML setting. But, as suggested above, I feel it would be a better paper in itself (if admittedly not necessarily for this venue) if the authors considered how it would perform for large M and more general guidance could be given to researchers on where to switch between algorithms.", "limitations": "", "rating": 10, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "This is a good study. Despite the field being well-studied, their approach is fairly original. It really sets aside what has typically been done and asks, for this context, precisely what does the algorithm need to do and on what hardware would it run. The consideration of the appropriate level of parallelism based on both GPU hardware constraints and problem size considerations is thoughtful, and they get to a nice answer. The paper is quite well written and was pleasantly easy to read.\n\nThe GNN test doesn’t have a ton of detail, but is well conducted and does a nice job of rounding out a clear ML story from an otherwise pure HPC-focused study.", "weaknesses": "I debate whether this is a better baseline comparison that PyTorch’s RadixSelect. As the authors suggest, it solves a somewhat more difficult problem – sorting. I don’t disagree with the comparison for the GNNs, as it is the operational default, but it was more difficult for me to interpret the value of the speedups is in the synthetic data top-k test. \n\nI would have liked the authors to push the bounds on the algorithm a little more. I understand the regime they tested is representative of a standard family of networks, but I would have liked to see another version of at least Figure 4 which looked at increasing M to ~10^4 or even 10^5 . This would indicate potential places where GPU cache/warp size limitations cause behavior to deteriorate, or suggest where the line between better performance with this algorithm vs more traditional approaches like RadiK is crossed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EzG5A7dOPb", "reviewer": "Reviewer_ZVep", "summary": "RTop-K is an efficient parallel row-wise top-k selection algorithm for GPUs, leveraging a binary search-based approach to significantly enhance performance in neural network training and inference. This algorithm achieves speedups ranging from 4.25 to 9.51 times compared to existing GPU top-k implementations, and it also improves the overall training workflow of MaxK-GNN by 9.76% to 31.53%. Furthermore, RTop-K, when combined with the Early Stopping technique, maintains model accuracy while enhancing efficiency through reduced resource requirements.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. RTop-K shows a 4.25- to 9.51-fold speedup over existing state-of-the-art GPU-based row-wise top-k implementations, demonstrating in particular performance improvements in GNN training and inference.  It is meaningful that RTop-K is designed based on efficient parallel processing and resource optimization.\n2. RTop-K adopts an approximation method to select the top k elements from each row without sorting, reducing computational time complexity and accelerating neural network training.", "weaknesses": "1. What is the method to overcome the limitations of approximation methods used by algorithms?\n2. Is resource optimization a general methodology? In the paper, the performance of RTop-K is likely to depend on a particular GPU architecture (A6000). Therefore, it is necessary to review whether usabilty is secured in other GPU environments or structures.\n3. RTop-K is optimized for limited-sized vectors, but there seems to be a lack of discussion on how resource consumption and memory management are performed on very large datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PwhcoSHHjr", "reviewer": "Reviewer_TSWj", "summary": "This paper presents RTop-K, a GPU-based algorithm that enhances the efficiency of row-wise top-K selection, specifically targeting Graph Neural Network (GNN) applications. The authors use a binary search strategy with optional early stopping to streamline top-K selection, offering a significant reduction in computational demands while preserving accuracy. Experiments show RTop-K provides speedups of 4.25× to 9.51× relative to PyTorch’s top-K implementation and achieves up to 31.53% faster training when integrated into GNN workflows.", "questions": "Except for the weaknesses above, I have several questions about this work.\n\n1, Just integrating the binary search into the top-K selection sounds not innovative, so can the authors dig more novelty from this work? Besides, I am curious that what searching performance improvement will using binary search bring, compared with not using binary searching. \n\n2, The authors claim the parallel top-K selection algorithm, but what are the specific 'parallel' mechanism is NOT clear. It seems that, just many single top-K selection are executed on GPU platform. Can the authors explain this? If my understanding is correct, then this paper should not emphasize the \"parallel\"?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "1,Relevance and Novelty: The paper addresses an important computational bottleneck in GNN acceleration, offering an innovative approach with practical speed improvements for real-world applications.\n\n2,Comprehensive Experimental Validation: The experimental analysis is thorough, showing RTop-K’s efficacy both as a standalone kernel and within the GNN training workflow across various architectures and datasets.\n\n3,Effective Use of Early Stopping: The early stopping mechanism is a valuable addition, providing computational savings without significantly impacting model accuracy.\n\n4,Reproducibility: The paper provides detailed descriptions of the algorithm and implementation, along with pseudocode, which improves replicability and clarity for the audience.", "weaknesses": "1,Ablation Study on Implementation Optimization: The paper describes a GPU implementation with several stages (loading, searching, and selecting) but could benefit from a breakdown of how each stage contributes to the overall speedup. A detailed analysis of performance gains per stage would shed light on optimization details. Additionally, an ablation study examining how binary search and early stopping impact efficiency would highlight their specific contributions and help readers better understand the trade-offs involved.\n\n2,Scalability Results on Hidden Dimensions: While the paper provides scalability results for different values of K, it lacks scalability analysis on various hidden dimensions, which is equally important for understanding the algorithm’s performance on high-dimensional data, as seen in some GNN and non-GNN applications.\n\n3,Guidance on Hyper-Parameter Tuning: The RTop-K algorithm involves critical hyper-parameters, specifically the precision value (ϵ') in Algorithm 1 and max_iter in Algorithm 2. Some discussion on how to set these parameters effectively across different tasks would enhance usability. A sensitivity analysis for these parameters could guide users in balancing performance with computational efficiency.\n\n4,Lack of Comprehensive Baselines: The paper compares RTop-K primarily with PyTorch’s row-wise top-K implementation. Including comparisons with other top-K selection algorithms, such as Avg-TopK and TopK-SGD, which are popular in neural network training, would provide a more complete evaluation of RTop-K’s performance and position it within the broader field.\n\n5,Discussion on Broader Applicability: While the paper focuses on GNN applications, top-K selection is widely used in various domains, such as expert selection in Mixture of Experts and vector search in large models like GPTs. A discussion on potential extensions of RTop-K to other contexts would strengthen the paper, showcasing its broader utility across different neural network architectures and tasks.\n\n6.Contribution. summary of something, in this case, GPU top-K selection alrorithms, should not be the contribution of one research paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)