RTop-K: Ultra-Fast Row-Wise Top-K Selection for Neural Network Acceleration on GPUs

Q4 - Compare with other top-K solutions

First, it is worth noting that Avg-TopK and TopK-SGD are applications of top-k in neural networks rather than GPU implementations.

Next, we want to make it clear that PyTorch's row-wise top-K implementation represents the state-of-the-art solution for the row-wise top-K scenario. Currently, other top-K solutions, including the reference [1] (Shanbhag et al.) provided by Reviewer 3azA, as well as those mentioned in our paper—Dr. Top-k (Gaihre et al.), AIR Top-k (Zhang et al.), and RadiK (Li et al.)—cannot be directly applied to row-wise top-K scenarios. If they must be applied, for example, in a row-wise top-K scenario with 65,536 rows, the kernel of the above solutions would need to be called 65,536 times to process the rows. The benchmark results are as follows (number of rows $N=65536$, vector size $M=[1024,2048]$):

The reference [1] (Shanbhag et al.) includes both RadixSelect and Bitonic kernels. It supports a minimum vector size $M = 1024$, but errors occur when the ratio of $k/M$ reaches a certain threshold.

As shown, the reference implementation is hundreds of times slower than both the PyTorch implementation and our implementation. This demonstrates that top-K implementations designed for single large vectors (typically with vector sizes in the millions) are not directly applicable to row-wise top-K scenarios (large batches of small vectors, which are common in deep learning and GNNs). Our innovations and contributions lie in designing and optimizing specifically for the row-wise top-K scenario.

[1] Anil Shanbhag et al. Efficient Top-K Query Processing on Massively Parallel Hardware.

Q5 - Discussion on broader applicability

Thanks for the suggestions. The RTop-K can indeed be used for Kth value search in Mixture of Expert (MoE) router network computation and vector search in the Retrieval Augmented Generation (RAG) application, especially for workloads with many short vectors to process. In the large batch training stage of LLMs [1] or the diffusion transformer models [2], with a batch size of 200 and a sequence length of 1000, there could easily be 0.2 million tokens to process for each Mixture of Experts layer. In such cases, the row-wise TopK operation could potentially become a bottleneck. We will add a section discussing these broader applications in the final version of the paper.

[1] Dai, D., Deng, C., Zhao, C., Xu, R. X., Gao, H., Chen, D., ... & Liang, W. (2024). Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. arXiv preprint arXiv:2401.06066.

[2] Sun, H., Lei, T., Zhang, B., Li, Y., Huang, H., Pang, R., ... & Du, N. (2024). EC-DIT: Scaling Diffusion Transformers with Adaptive Expert-Choice Routing. arXiv preprint arXiv:2410.02098.

Q6 - Contribution

Our contribution lies in the following key aspects:

1. Optimizing the row-wise TopK kernel to achieve performance several times faster than PyTorch.
2. Integrating RTopK into GNNs training systems, resulting in up to a 1.3x overall speedup.

Optimizing the row-wise TopK kernel itself is a non-trivial research problem, as row-wise TopK is challenging to accelerate and can become a bottleneck in GNN training (as shown in Table 4, where it accounts for 12%-27% of total training time).

We respectfully disagree with the notion that GPU top-K selection algorithms themselves should not be the contribution of a research paper. GPU top-K selection is a highly impactful research area that supports numerous applications. In fact, there are already multiple research papers [1][2][3] dedicated to optimizing GPU top-K selection algorithms.

[1] Anil Gaihre et al. Dr. Top-k: Delegate-centric Top-k on GPUs. SC '21.
[2] Jingrong Zhang et al. Parallel top-k algorithms on GPU: A comprehensive study and new methods. SC ’23.
[3] Yifei Li et al. RadiK: Scalable and optimized GPU-parallel radix top-k selection. ICS ’24.

Thank you for taking the time to read and review our paper!

Q1 - A breakdown of each stage

We use the clock64() method within the kernel to roughly measure the clock cycles occupied by each stage and calculate their relative proportions.

It can be observed that as max_iter increases, the time spent on the loading and selecting stages remains relatively constant, while the time spent on the searching stage continues to grow. This aligns with our design expectations. However, it is important to note that measuring clock cycles within the kernel using clock64() is limited to individual threads, and the results represent an average across all threads. Additionally, the total clock cycles measured within the kernel translate to a runtime that is shorter than the standard timing test we use for overall kernel execution. This is because the standard timing test (used for comparison with PyTorch) includes kernel on/off-chip overhead and some synchronization costs.
The implementation and optimization effectiveness of the kernel depend on our skillful use and deep understanding of warp-level primitives, as well as efficient management of threads, warps, blocks, registers, and shared memory. In this context, the measurement of clock cycles within the kernel serves primarily as a validation of the optimization effectiveness.

Q2 - Scalability results on hidden dimensions

In the newly added Appendix A, we successfully derived the expected number of iterations for RTopK's no early stopping version (Algorithm 1):

$$E(n) = \log_2 \left( 2M \sqrt{\frac{\ln M}{\pi}} \right) - \frac{1}{2 \ln 2} \left( \Phi^{-1}\left( 1 - \frac{k}{M} \right) \right)^2$$

According to our analysis, the time complexity of Algorithm 1 for single-vector top-k selection is less than $O(M\log M)$ and slightly greater than $O(M)$. This establishes the relationship between the time complexity and the hidden dimension $M$. For the overall row-wise top-k operation, which includes $N$ rows, the total time complexity ranges from $O(NM)$ to $O(NM\log M)$.

In the newly added Appendix B, we performed an in-depth analysis of RTopK's performance under different values of $M$. For the no early stopping version:

• When $M$ is below 3072, the speed of RTop-K is generally more than twice that of PyTorch.
• When $M$ is between 3072 and 6144, RTop-K is still generally faster than PyTorch.
• For $M > 8192$, our current shared memory-based acceleration strategy cannot be directly applied on the A6000 GPU due to shared memory limitations.

From these findings, we conclude that RTopK holds a clear advantage in common neural network scenarios.

Q3 - Discussion on precision and max_iter

First, please allow us to share our intriguing discovery: in the no early stopping version, the choice of precision has almost no impact on speed, even when set to maximum precision for floating-point type ($\epsilon = 0$). When $\epsilon = 0$, RTopK performs as an accurate top-k selection and is still several times faster than PyTorch in common neural network scenarios. For more details, please refer to the latter part of Appendix B.

Therefore, the optimal precision choice for Algorithm 1 is simply $\epsilon = 0$.

For Algorithm 2 (early stopping version), it can be observed from Figure 5 that early stopping has virtually no impact on accuracy; in some cases, early stopping even achieves higher accuracy than no early stopping (possibly due to the additional non-linearity introduced by early stopping). Based on the experimental results, setting max_iter to 4 or 5 strikes a good balance between acceleration and accuracy. A more detailed analysis would depend on additional practical applications of the row-wise Top-K layer.

Q7 - Advantages of binary search algorithm

The advantages of binary search-based top-k selection lie in its simplicity and resource efficiency, making it more effective in row-wise top-k scenarios and even several times faster than PyTorch's official implementation.

Here, we compare with two traditional algorithms, RadixSelect and Bitonic, as presented in the reference [1] (Shanbhag et al.) provided by Reviewer 3azA:

• RadixSelect
  Refer to the implementation: radixSelectTopK.cuh.
  RadixSelect involves a series of steps, including initialization, histogram construction, indexing, sorting and selection, and extraction of the top-K elements.

• Bitonic
  Refer to the implementation: bitonicTopK.cuh.
  Bitonic sorting requires multiple rounds of extensive element exchanges and includes memory read-write dependencies.

In contrast, binary search-based top-k selection, during its main stage (the searching stage), involves only parallel counting operations. These operations can efficiently utilize warp-level primitives for reduction, with no element movement or memory read-write dependencies. This makes binary search-based top-k selection significantly faster for small vectors compared to traditional algorithms.

However, due to time complexity considerations, as the vector size grows beyond a certain point, RTopK may lag behind traditional algorithms. Nonetheless, RTopK retains its advantage in common neural network scenarios.

[1] Anil Shanbhag, Holger Pirk, and Samuel Madden. Efficient Top-K Query Processing on Massively Parallel Hardware.

Q8 - About "parallel"

The parallelism of the RTopK implementation for the row-wise top-K scenario is reflected on two levels:

1. Intra-warp parallelism:
   RTopK uses a single warp to process one row (a limited vector). During this process, the loading, searching, and selecting stages involve collaboration among all threads within the warp. Specifically, operations such as tree-reduction and Ballot rely on warp-level primitives to achieve efficient parallel processing within the warp.

2. Inter-row parallelism:
   The RTopK kernel processes all rows of the matrix in parallel. Since the computation and memory resources allocated to each row's processing are mutually independent, the parallelism at this level is relatively straightforward. However, we still need to carefully allocate resources to each block to maximize the efficiency of inter-row parallelism. As mentioned in Appendix B, we allocate ⌊8192/M⌋ warps per block, which has proven to be an effective strategy.

While our primary focus is on designing and optimizing the algorithms, parallelism is also an important aspect of our design.

Dear Reviewer,

Thank you once again for your valuable feedback, we have carefully addressed your concerns and incorporated your suggestions. In the rebuttal, we added many other ablation and comparative studies ensuring that RTopK is thoroughly studied. Furthermore, we have open-sourced our work (currently in anonymous github).

As the discussion phase is nearing its end, we would like to kindly follow up to confirm whether our response has sufficiently addressed your concerns. If so, we would greatly appreciate it if you could consider raising the score further to 6: marginally above the acceptance threshold. We firmly believe that our work has the potential to contribute significantly to the field and meets the criteria for acceptance. If fortunate enough to receive the reviewers' recognition, we are committed to contributing this work to the official PyTorch library to benefit the broader community.

Please do not hesitate to let us know if you have any additional questions or concerns; we would be happy to address them.

Thank you again for your time and thoughtful input.

Thank you for taking the time to read and review our paper!

Q1 - Overcome the limitations of approximation

First, please allow us to share our intriguing discovery: in the no early stopping version, the choice of precision has almost no impact on speed, even when set to maximum precision for floating-point type ($\epsilon = 0$). When $\epsilon = 0$, RTopK performs as an accurate top-k selection and is still several times faster than PyTorch in common neural network scenarios.

For more details, please refer to the newly added Appendix A and B. This significantly boosts our confidence in contributing our work to PyTorch's official library, though it will still demand considerable engineering efforts.

To completely overcome the limitations of approximation, using the accurate version provides the optimal experience!

Of course, the early stopping version is also a viable choice. In Table 2, we present theoretical simulation tests for the relative error and hit rate of the approximation algorithm. The analysis shows that the relative error of the approximation algorithm can easily be controlled within 5%. Tests on the impact of this approach on actual test accuracy are shown in Figure 5 (blue dots represent the accuracy with early stopping, and the green horizontal line represents the accuracy without early stopping). It can be observed that early stopping has almost no impact on accuracy. Therefore, when applying the row-wise top-k layer in typical models, early stopping can be used to achieve greater speedup with minimal impact on accuracy.

Q2 - Testing on other GPUs

As shown in Appendix B, RTopK generally achieves speeds more than twice that of PyTorch when the vector size $M < 3072$. This is sufficient to cover the vast majority of row-wise top-k use cases in neural networks and is applicable to other mainstream GPUs. We extended Table 3 to include tests on A100 and Quadro RTX 6000 GPUs. The results are as follows:

A100:

Quadro RTX 6000:

It can be observed that RTopK outperforms PyTorch on both A100 and Quadro RTX 6000 GPUs.

Q3 - Resource and memory analysis on large datasets

RTopK is a fast operator for computing row-wise top-k. In its application scenarios, while the row-wise top-k operation has a certain computational complexity, it is typically not relatively costly in terms of resource and memory usage.

For example, in MaxK-GNN, the SpMM operation is the primary consumer of memory, whereas the memory usage of RTopK is comparatively minimal. Consider the Reddit graph, which has 232,965 nodes and 114,615,891 edges. With a hidden dimension of 256, the total memory traffic for a single SpMM operation is 13.13 GB, while the total memory traffic for a single row-wise top-k operation is only 0.29 GB.

Thus, our focus is primarily on optimizing the runtime of RTopK, as its memory resource requirements will not become a bottleneck.

Dear Reviewer,

We hope our response has addressed all of your concerns. Please feel free to let us know if you have any additional questions or feedback. This project is already open-sourced (currently in anonymous github), and we assure you that we will do our utmost to contribute this work to the official PyTorch library. If you could kindly consider raising your evaluation, your support would greatly help accelerate this process. Thank you very much for your time and thoughtful consideration.

We would like to further clarify question 3.

Q3 - Resource and memory analysis on large datasets

RTop-K is designed for the row-wise TopK operator over a large number of small vectors, which can already represent large datasets. For instance, the Ogbn-products dataset used in the experiment (Table 4/Figure 5 in the paper) contains $N = 2,449,029$ rows. With an embedding dimension of $M = 1024$, the embedding matrix (input matrix in the row-wise TopK scenario) comprises 2,507,805,696 elements, occupying 10.03 GB of memory. This already qualifies as a large dataset.

Additionally, for the large number of rows $N$ with large vector size $M$, we have added profiling results up to $N = 65536$, $M \leq 8192$. For details, please see Figure 6 in Appendix B. With the early stop algorithm, RTop-K kernel consistently outperforms Pytorch kernel.

For memory resource usage, we can classify the memory resource usage into two types: global memory allocation and global memory traffic (read/write).

For global memory allocation, RTop-K kernel requires no additional global memory allocation. The memory usage of RTop-K is even lower than that of conventional TopK Pytorch kernel, as it will only require necessary memory allocation for input matrix (batch of vectors) and output matrix (batch of vectors).

In terms of global memory traffic (read/write) during the kernel computation, RTop-K consumes significantly less memory traffic compared to other operators in the neural network workflow. It involves only $N \times M$ global memory reads and $N \times k$ global memory writes, which represent a memory traffic similar to that of the element-wise operator, such as ReLU. All other operations are completed within shared memory.

For example, in the Reddit graph, which has 232,965 nodes and 114,615,891 edges. With a hidden dimension of 256, the total global memory traffic (read/write) for a single row-wise top-k operation is only 0.29 GB (memory traffic, not additional static memory usage). While in MaxK-GNN, other operators such as SpMM, a single SpMM forward consumes totally 13.13 GB memory traffic(read/write).

Thus, our focus is primarily on optimizing the runtime of RTop-K, as the kernel itself doesn’t require additional global memory allocation other than necessary input/output memory, and it requires very minimum global memory traffic usage, so the memory resource will not become a bottleneck.

Q6 - GNN & RTop-K training time breakdown

Below is the breakdown of training time for GCN & RTop-K models on several graph datasets (with top-k acceleration, hidden dimension $M = 256$, $k = 32$):

Is TopK a bottleneck in the GNN?

The row-wise TopK operation serves as a nonlinear component in our GNN framework, introducing structural sparsity and improving performance. However, without a specialized design, the row-wise TopK itself can become a bottleneck during GNN training and inference. To mitigate this, we optimize the row-wise TopK kernel using the RTop-K implementation, achieving further acceleration.

Q3 - Theoretical analysis

Thank you for your suggestion to include additional theoretical derivations!

For Algorithm 1 (no early stopping version), we conducted mathematical derivations and more detailed theoretical simulations. We successfully derived an expression for the expected number of iterations:

$$E(n) = \log_2 \left( 2M \sqrt{\frac{\ln M}{\pi}} \right) - \frac{1}{2 \ln 2} \left( \Phi^{-1}\left( 1 - \frac{k}{M} \right) \right)^2$$

This expression matches the simulation results well, as detailed in the newly added Appendix A.

One surprising finding is that setting the precision of the no early stopping version to maximum for floating-point type ($\epsilon=0$) has almost no impact on speed. When $M < 3072$, the speed of RTop-K ($\epsilon=0$) is generally more than twice that of PyTorch. Detailed results can be found in the newly added Appendix B. This significantly boosts our confidence in contributing our work to PyTorch's official library, though it will still demand considerable engineering efforts.

For Algorithm 2 (early stopping version), since the specific execution process involves a series of discrete decisions, a more in-depth theoretical analysis was supported through simulation experiments. In Table 2, we conducted simulation tests on the relative error and hit rate of early stopping. The analysis shows that the relative error of early stopping can be easily controlled within 5%. Figure 5 illustrates the impact of early stopping on the actual test accuracy of the model (blue dots represent the accuracy with early stopping, while the green horizontal line represents the accuracy without early stopping). The conclusion is that early stopping has almost no impact on accuracy; in some cases, the accuracy with early stopping is even higher than without it. Therefore, early stopping is well-suited for further accelerating the application of this algorithm in neural networks.

Q4 - Mark speedup on Fig. 4

Got it, we will highlight the acceleration effect in Figure 4 of the final version of the paper.

Q5 - Nsight profiling results

Below are the NSight Compute profiling results of RTopK (no early stopping, $\epsilon=0$) and PyTorch TopK, tested on an A6000 GPU with $N=65536$, $M=768$, and $k=64$:

As shown, the RTopK kernel achieves very high Compute Throughput and Memory Throughput, with both being nearly balanced, fully leveraging the GPU's performance.

In contrast, PyTorch's TopK implementation is split into multiple kernels, primarily consisting of several calls to radixFindKthValues and computeBlockwiseWithinKCounts, followed by a single gatherTopK call. None of these kernels fully utilize Compute Throughput and Memory Throughput.

Thank you for taking the time to read and review our paper!

Q1 - Compare with other top-K solutions

We want to make it clear that PyTorch's row-wise top-K implementation represents the state-of-the-art solution for the row-wise top-K scenario. Currently, other top-K solutions, including the reference you provided (Shanbhag et al.), as well as those mentioned in our paper—Dr. Top-k (Gaihre et al.), AIR Top-k (Zhang et al.), and RadiK (Li et al.)—cannot be directly applied to row-wise top-K scenarios. If they must be applied, for example, in a row-wise top-K scenario with 65,536 rows, the reference kernel would need to be called 65,536 times to process the rows. The benchmark results are as follows (number of rows $N=65536$, vector size $M=[1024,2048]$):

The reference (Shanbhag et al.) includes both RadixSelect and Bitonic kernels. It supports a minimum vector size $M = 1024$, but errors occur when the ratio of $k/M$ reaches a certain threshold.

As shown, the reference implementation is hundreds of times slower than both the PyTorch implementation and our implementation. This demonstrates that top-K implementations designed for single large vectors (typically with vector sizes in the millions) are not directly applicable to row-wise top-K scenarios (large batches of small vectors, which are common in deep learning and GNNs). Our innovations and contributions lie in designing and optimizing specifically for the row-wise top-K scenario.

In the table, we also compare PyTorch-Sort and PyTorch-Topk. Result shows that PyTorch-Topk is faster than PyTorch-Sort in most cases, and is only slightly slower than PyTorch-Sort in only 1 case when $k/M$ is too large (k=512 M=1024). In both of those cases, our method RTop-K shows a significant advantage over both PyTorch-Sort and PyTorch-Topk.

Q2 - Applicability

First, RTopK indeed has more application scenarios. For example, it can be used for Kth value search in Mixture of Expert (MoE) router network computation and vector search in the Retrieval Augmented Generation (RAG) application, especially for workloads with many short vectors to process. In the large batch training stage of LLMs [1] or the diffusion transformer models [2], with a batch size of 200 and a sequence length of 1000, there could easily be 0.2 million tokens to process for each Mixture of Experts layer. In such cases, the row-wise TopK operation could potentially become a bottleneck.

Second, the row-wise TopK layer can generally be used as an non-linear activation layer in neural network models (currently used in our GNNs), similar to ReLU nonlinearity. Compared to ReLU, it can output a row-balanced structural sparsity pattern, which enables efficient utilization of this sparsity in applications like MaxK-GNN — a significant and successful example. On the other hand, it can provide a novel non-linearity, potentially leading to better theoretical accuracy for DNNs in certain scenarios.

Thank you for your insight! We will add a section discussing these applications in the final version of the paper.

[1] Dai, D., Deng, C., Zhao, C., Xu, R. X., Gao, H., Chen, D., ... & Liang, W. (2024). Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. arXiv preprint arXiv:2401.06066.
[2] Sun, H., Lei, T., Zhang, B., Li, Y., Huang, H., Pang, R., ... & Du, N. (2024). EC-DIT: Scaling Diffusion Transformers with Adaptive Expert-Choice Routing. arXiv preprint arXiv:2410.02098.

Thank you very much for your recognition of our work. Your suggestions are highly reasonable, as traditional algorithms like RadixSelect can outperform our method when the vector size $M$ grows beyond a certain point. Determining the critical switching point is indeed necessary.

Please refer to the newly added Appendix A, where we successfully derive the expression for the expected number of iterations in Algorithm 1 (no early stopping version):

$$E(n) = \log_2 \left( 2M \sqrt{\frac{\ln M}{\pi}} \right) - \frac{1}{2 \ln 2} \left( \Phi^{-1}\left( 1 - \frac{k}{M} \right) \right)^2$$

According to our analysis, the time complexity of Algorithm 1 is less than $O(M\log M)$ and slightly greater than $O(M)$. In comparison, the underlying RadixSelect in PyTorch has a time complexity of $O(M)$. Of course, this is purely theoretical. The actual situation is more complex. Please refer to Appendix B. For the no early stopping version, when $M$ is below 3072, the speed of RTop-K is generally more than twice that of PyTorch. When $M$ is between 3072 and 6144, RTop-K is generally faster than PyTorch. (For $M > 8192$, our current shared memory-based acceleration strategy cannot be directly applied on the A6000 GPU due to shared memory limitations.)

We are also delighted to share: the choice of precision has almost no impact on the speed of the no early stopping version, even when set to maximum precision for floating-point type ($\epsilon = 0$). When $\epsilon = 0$, RTopK performs as an accurate top-k selection and is still several times faster than PyTorch in common neural network scenarios. This significantly boosts our confidence in contributing our work to PyTorch's official library, though it will still demand considerable engineering efforts.

We also notice that when $M$ is a multiple of 256 or 512, the performance of RTopK drops. This is likely due to shared memory bank conflicts, and we are confident in addressing this issue in the final version of the paper (possibly by adding padding in the data access structure to break the conflict pattern).

Thank you again for recognizing our work! May we ask for your even stronger approval? :)