We introduce FastGraph, a novel GPU-optimized k-nearest neighbor algorithm specifically designed to accelerate graph construction in low-dimensional spaces (2-10 dimensions), critical for high-performance graph neural networks. Our method employs a GPU-resident, bin-partitioned approach with full gradient-flow support and adaptive parameter tuning, significantly enhancing both computational and memory efficiency. Benchmarking demonstrates that FastGraph achieves a 20-40x speedup over state-of-the-art libraries such as FAISS, ANNOY, and SCANN in dimensions less than 10 with virtually no memory overhead. These improvements directly translate into substantial performance gains for GNN-based workflows, particularly benefiting computationally intensive applications in low dimensions such as particle clustering in high-energy physics, visual object tracking, and graph clustering.
@article{arxiv.2511.10442,
title = {FastGraph: Optimized GPU-Enabled Algorithms for Fast Graph Building and Message Passing},
author = {Aarush Agarwal and Raymond He and Jan Kieseler and Matteo Cremonesi and Shah Rukh Qasim},
journal= {arXiv preprint arXiv:2511.10442},
year = {2025}
}