FastGraph: Optimized GPU-Enabled Algorithms for Fast Graph Building and Message Passing

TL;DR

FastGraph is a GPU-optimized algorithm that significantly accelerates graph construction and message passing in low-dimensional spaces, enabling faster GNN workflows with minimal memory overhead.

Contribution

It introduces a novel GPU-resident, bin-partitioned k-nearest neighbor algorithm with adaptive tuning, outperforming existing libraries in speed and efficiency.

Findings

Achieves 20-40x speedup over FAISS, ANNOY, SCANN

Operates efficiently in dimensions less than 10

Enhances GNN performance in applications like physics and image analysis

Abstract

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.