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

A team of researchers from UC Berkeley and UT Austin have published Flash-KMeans, an open source implementation of vanilla Lloyd’s k-means clustering on GPUs that they claim is more than 200x faster than FAISS without any approximations. How? By rearranging the way data is moved around instead of approximating. Flash-KMeans is available under the Apache 2.0 license and can be installed via ‘pip install flash-kmeans’.

On an NVIDIA H200 GPU in FP16 with 128 dimensions, Flash-KMeans achieves impressive results including being 33x faster than NVIDIA’s own cuML, up to 17.9x faster end-to-end compared to state-of-the-art baselines, and up to 10.5x faster than fastkmeans in an out-of-core setting with 1 billion points. Furthermore, the two custom kernels achieve up to 21.2x and 6.3x speedups which contribute to these end-to-end improvements.

1How it gets the speedup

To achieve these gains, Flash-KMeans uses two clever memory optimizations. First, FlashAssign combines distance calculation and online argmin to avoid storing the entire NxK distance matrix in memory, reducing the IO complexity from O(NK) to O(Nd+Kd). Second, Sort-Inverse Update sorts points according to their cluster IDs, turning atomic writes into segment reductions and reducing atomic contention. Importantly, these optimizations don’t affect the final clusters, only how much data needs to be moved between different levels of the GPU’s memory.

2Why it matters

K-means is a fundamental building block in many modern AI applications such as vector search indexing, sparse attention routing, KV cache compression, low-bit KV quantization, and diffusion transformers. In these workflows, clustering is frequently used within training and inference loops, but what matters most is not theoretical complexity but rather latency per invocation. With a simple drop-in replacement for faster exact clustering, you can get the same results in less time and money without having to worry about sacrificing accuracy.