Research map / Architectures and efficiency / RNNs and normalization
RTop-K: Ultra-Fast Row-Wise Top-K Selection for Neural Network Acceleration on GPUs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes RTop-K, a warp-per-row binary-search top-k kernel for GPUs with optional early stopping, showing large kernel speedups over PyTorch and modest end-to-end MaxK-GNN training speedups.
Problem
Row-wise top-k selection over very large batches of short vectors (hidden dim up to ~1024) is inefficient on GPUs, since existing GPU top-k implementations target a few large vectors and PyTorch's RadixSelect-based top-k is overly complex and sorts its outputs.
Approach
A binary search over the value range finds a threshold that selects k elements, with a second pass using the minimum to fill borderline ties; warp-level primitives (shuffle, ballot, popcnt) implement one warp per row, and capping the number of search iterations (early stopping) trades a small approximation for speed.
Key results
- Average kernel speedup over PyTorch of 11.49x with early stopping (max_iter=2) and 7.29x without early stopping, for M=256 across N and k.
- MaxK-GNN end-to-end training speedups of 11.97% to 12.21% on Reddit, 32.48% to 33.29% on Flickr, 22.00% to 22.74% on Ogbn-products, and 31.21% to 32.42% on Yelp.
- With max_iter>=5 (k>=32), early-stopping top-k keeps max and min relative errors at or below 5%; for k>=64, 4 iterations give hit rates above 80%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add early stopping with a small max_iter (ablation) | RTop-K row-wise top-k kernel · Row-wise top-k kernel, average across tested settings | speedup vs PyTorch: 7.29 -> 11.49 | Helped |
| Add early stopping to row-wise top-k selection (ablation) | MaxK-GNN · MaxK-GNN testing accuracy across models and graphs | testing accuracy | No clear effect |
| Use RTop-K instead of PyTorch row-wise top-k | MaxK-GNN · Reddit, Flickr, Ogbn-products, and Yelp | overall training speedup: 11.97% to 33.29% | Helped |
| Vary precision epsilon (ablation) | RTop-K without early stopping · Row-wise top-k kernel | speed | No clear effect |
| Increase vector size M (ablation) | RTop-K without early stopping · Row-wise top-k kernel, M between 6144 and 8192 | speed vs PyTorch | Hurt |
Surprising findings
- Precision (epsilon) has almost no impact on kernel speed, and exact selection (epsilon=0) exits within 16 iterations in most cases.
- Early stopping sometimes yields higher test accuracy than exact top-k in MaxK-GNN training.
Limitations
- Speedup shrinks with larger row length M and the design is slower than PyTorch beyond M about 6144-8192 due to O(M log M) complexity and shared memory limits (M > 8192 unsupported).
- Early-stopping accuracy analysis relies on normally distributed synthetic vectors, and GIN on Flickr showed accuracy degradation.
- Overall training speedups are measured only for MaxK-GNN with a single configuration (M=256, k=32).
Experimental setup
Datasets: Flickr Yelp Reddit Ogbn-products
Benchmarks: Flickr Yelp Reddit Ogbn-products
Metrics: kernel latency (ms) speedup over PyTorch overall training speedup test accuracy hit rate relative error
Baselines: PyTorch topk (RadixSelect) Dr. Top-k AIR Top-k RadiK GridSelect
Builds on: MaxK-GNN PyTorch topk Dr. Top-k AIR Top-k RadiK
Closest papers by meaning
- TPU-KNN: K Nearest Neighbor Search at Peak FLOP/s (NeurIPS 2022 · similarity 0.85)
- ST$_k$: A Scalable Module for Solving Top-k Problems (NeurIPS 2024 · similarity 0.83)
- Efficient Use of Limited-Memory Accelerators for Linear Learning on Heterogeneous Systems (NeurIPS 2017 · similarity 0.83)
- FlashRNN: I/O-Aware Optimization of Traditional RNNs on modern hardware (ICLR 2025 · similarity 0.82)
- Poplar: Efficient Scaling of Distributed DNN Training on Heterogeneous GPU Clusters (AAAI 2025 · similarity 0.82)
- FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (NeurIPS 2024 · similarity 0.82)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.