Flash-KMeans is an open‑source library that speeds up exact Lloyd’s k‑means on GPUs without changing the math. Modern AI pipelines now call k‑many times during training and inference, so latency per call matters more than raw FLOPs. Flash‑KMeans tackles the two memory‑bound stages of the algorithm.
The first bottleneck is the assignment step, where a full N×K distance matrix is built in high‑bandwidth memory, written, then read back for argmin. Flash‑KMeans replaces this with FlashAssign, which streams tiles of points and centroids into on‑chip SRAM, fuses distance computation with an online argmin, and never materializes the full matrix. This cuts the IO complexity from O(NK) to O(Nd+Kd) and yields up to 21.2× speedup on the assignment kernel.
The second bottleneck is the centroid update, where standard scatter‑style atomic adds cause contention as many threads write to the same hot centroid. Flash‑KMeans uses Sort‑Inverse Update: it sorts the assignment vector by cluster ID, turns scattered writes into contiguous segment reductions, and performs a single atomic add per segment. This reduces atomic overhead and gives up to 6.3× speedup on the update kernel.
On an NVIDIA H200 with FP16 data, Flash‑KMeans reports up to 17.9× end‑to‑end speedup over the best baseline, 33× over NVIDIA cuML, and over 200× over FAISS. It also scales out‑of‑core to one billion points, finishing an iteration in 41.4 seconds versus 261.8 seconds for a baseline, and reduces tuning overhead by up to 175×.
The API mirrors FAISS and scikit‑learn. A simple call clusters a batched tensor:
import torch
from flash_kmeans import batch_kmeans_Euclid
x = torch.randn(32, 75600, 128, device=”cuda”, dtype=torch.float16)
ids, centers, _ = batch_kmeans_Euclid(x, n_clusters=1000, tol=1e-4, verbose=True)
A scikit‑learn‑style FlashKMeans class is also available, auto‑dispatching kernels based on shape and dtype, supporting multi‑GPU out‑of‑core execution.
Flash‑KMeans delivers exact k‑means results at a fraction of the latency, enabling online vector‑search indexing, attention routing, KV‑cache compression, low‑bit quantization, and diffusion transformer workloads that previously relied on offline preprocessing.
#AI #Product #GPU #MachineLearning #KMeans #FlashKMeans