k-means has been an offline tool for decades. You run it once to preprocess data, then move on. A team of researchers from UC Berkeley and UT Austin released Flash-KMeans, a new open-source library that targets a different setting. Modern AI pipelines now call k-means inside training and inference loops. At that frequency, latency per call matters more than theoretical FLOPs.<\/p>\n
Flash-KMeans is an IO-aware implementation of standard Lloyd\u2019s k-means. It does not change the math, and it does not approximate. It only restructures how the algorithm moves data on a GPU. On an NVIDIA H200, the research team reported up to 17.9\u00d7 end-to-end speedup over the best baseline. Against NVIDIA cuML they report 33\u00d7. Against FAISS they report over 200\u00d7. <\/p>\n
Flash-KMeans is a batched k-means library written in Triton GPU kernels. It ships under Apache 2.0 and installs with The output is mathematically identical to standard Lloyd\u2019s k-means. The speedup comes from kernel-level dataflow, not from skipping work. That separates it from algorithmic methods like triangle-inequality pruning or coreset sampling.<\/p>\n A standard Lloyd iteration has two stages. The assignment stage computes each point\u2019s distance to every centroid, then picks the nearest. The update stage averages the points in each cluster to form new centroids. Both stages are simple arithmetic. On GPUs, both are bottlenecked by memory, not compute.<\/p>\n The first bottleneck<\/strong> is the assignment stage. Standard code builds a full distance matrix D of shape N\u00d7K in High Bandwidth Memory (HBM). It writes the matrix, then reads it back to run argmin. For N=65536, K=1024, d=128, B=32, the distance math takes 2.6ms. Writing and consuming D takes about 23ms. The matrix is the cost, not the arithmetic.<\/p>\n Flash-KMeans replaces this with FlashAssign. The design borrows from FlashAttention. FlashAssign streams tiles of points and centroids from HBM into on-chip SRAM. It fuses distance computation with an online argmin. The full N\u00d7K matrix is never materialized. This cuts the dominant IO complexity from O(NK) to O(Nd + Kd). At the kernel level, FlashAssign reaches up to 21.2\u00d7. In one case it cut assignment from 122.5ms to 5.8ms.<\/p>\n The second bottleneck<\/strong> is the centroid update stage. Standard code uses scatter-style atomic adds. Each thread adds its point into a shared sum buffer keyed by cluster id. Many threads hit the same \u2018hot\u2019 cluster at once. That causes atomic contention and hardware serialization. The research team measured only 50 GB\/s effective bandwidth here on an H200.<\/p>\n Flash-KMeans replaces this with Sort-Inverse Update. It sorts the 1D assignment vector by cluster id using argsort. Identical cluster ids then form contiguous segments. Each thread block reduces a segment on-chip, then issues one atomic add per segment. The heavy point matrix is never physically permuted. Atomic operations drop from (O((K+NBN)d))(O((K + frac{N}{B_N})d)) . The kernel reaches up to 6.3\u00d7.<\/p>\n The research team test it on an H200 with CUDA 12.8, FP16 data, and d=128. They sweep N, K, and batch size B. They compare against four optimized baselines: fast_pytorch_kmeans, fastkmeans, cuML, and FAISS.<\/p>\n One failure mode matters for context. Standard PyTorch implementations run out of memory in large-K regimes. They cannot materialize the N\u00d7K matrix. FAISS is the industry-standard library under many production vector-search systems.<\/p>\n The library also runs out-of-core. On one billion points (K=32768, d=128), it finishes an iteration in 41.4s, against 261.8s for the baseline. It uses chunked stream overlap to hide PCIe transfer behind compute. A cache-aware compile heuristic also cuts tuning overhead by up to 175\u00d7, within 0.3% of tuned speed.<\/p>\n <\/p>\n Same Lloyd\u2019s math as standard k-means \u2014 faster only because of dataflow. Run clustering live, watch the update bottleneck, and size the IO it removes.<\/p>\n<\/div>\n <\/p>\npip install flash-kmeans<\/code>. <\/p>\nThe Two Bottlenecks It Attacks<\/strong><\/h2>\n
Benchmark<\/strong><\/h2>\n
\n\n
\n \nComparison<\/th>\n Reported speedup<\/th>\n Workload context<\/th>\n<\/tr>\n<\/thead>\n \n End-to-end vs best baseline<\/td>\n up to 17.9\u00d7<\/td>\n N=8M, K=1024 (large N, small K)<\/td>\n<\/tr>\n \n vs NVIDIA cuML<\/td>\n 33\u00d7<\/td>\n industry library<\/td>\n<\/tr>\n \n vs FAISS<\/td>\n over 200\u00d7<\/td>\n industry library<\/td>\n<\/tr>\n \n FlashAssign kernel<\/td>\n up to 21.2\u00d7<\/td>\n N=1M, K=8192 (assignment)<\/td>\n<\/tr>\n \n Sort-Inverse Update kernel<\/td>\n up to 6.3\u00d7<\/td>\n N=33M, K=4096 (update)<\/td>\n<\/tr>\n \n Out-of-core, large scale<\/td>\n up to 10.5\u00d7<\/td>\n N=400M, K=16384 vs fastkmeans<\/td>\n<\/tr>\n<\/tbody>\n<\/table>\n<\/figure>\n MTP Interactive Explainer<\/strong><\/h2>\n
Flash-KMeans: exact k-means, rebuilt around GPU memory<\/h2>\n