A team of researchers from UC Berkeley, MIT, UC Irvine, and UT Austin has open-sourced FlashLib, a GPU kernel library that accelerates classic machine learning algorithms — think KMeans, KNN, and PCA — by up to 208 times compared to NVIDIA's cuML 25.10. Built on Triton and CuteDSL, FlashLib covers 15 high-level operators and is designed to supercharge ML workflows and agent-based systems.
On an H200 GPU, Flash-KMeans hits 61% of peak FLOPs efficiency, while Flash-KNN achieves 85.2% of HBM memory bandwidth. Against cuML 25.10, FlashLib delivers a 26x speedup on KMeans clustering, 19x on KNN search, and 40x on HDBSCAN clustering. In TruncatedSVD decomposition, it reaches a staggering 208x peak acceleration under specific tolerance and algorithmic trade-offs. PCA, exact t-SNE, and MultinomialNB see speedups of 47x, 147x, and 49x, respectively.
As AI moves toward agentic workflows, classic ML operators are shifting from offline batch processing to online, real-time primitives. In long-chain reasoning tasks like scientific computing agents, models frequently call clustering, search, and dimensionality reduction operators, turning traditional offline kernels into system bottlenecks. FlashLib introduces a performance prediction API that can estimate a workflow's runtime and memory overhead with just about 5 microseconds of CPU time — no GPU profiling required — offering low-cost analysis for task planning and agent decision-making. FlashLib is now available on GitHub.