Flash-KMeans: Exact K-Means 200x Faster on GPUs

Researchers from UC Berkeley and UT Austin have released Flash-KMeans, an open-source GPU implementation of standard Lloyd's k-means clustering that reports running over 200x faster than FAISS. The notable part: it 'does not change the math, and it does not approximate. It only restructures how the algorithm moves data on a GPU.' It ships under Apache 2.0 and installs with a single 'pip install flash-kmeans.'

The benchmarks, measured on an NVIDIA H200 in FP16 at 128 dimensions, are steep: 33x faster than NVIDIA's own cuML, up to 17.9x end-to-end versus the best baseline, and up to 10.5x faster than fastkmeans on an out-of-core run over 1 billion points. The two custom kernels post their own gains — up to 21.2x and 6.3x — that combine into those end-to-end numbers.

How it gets the speedup

Two memory tricks do the work. FlashAssign fuses the distance computation with the online argmin, so the full N×K distance matrix is never written out to memory — dropping IO complexity from O(NK) to O(Nd+Kd). Sort-Inverse Update then sorts points by cluster ID, converting scattered atomic writes into contiguous segment reductions and cutting atomic contention. Neither step touches the result: you get identical clusters, just far less data shuffled across the GPU's memory hierarchy.

Why it matters

K-means is quiet plumbing underneath a lot of modern AI: vector-search indexing, sparse-attention routing, KV-cache compression, low-bit KV quantization, and diffusion transformers. These pipelines call clustering over and over inside training and inference loops, where the cost that actually bites is latency per call, not textbook big-O. A drop-in exact speedup means the same answers for less time and money — and because it's exact, there's no accuracy trade-off to argue about before adopting it.