# Flash-KMeans: Fast and Memory-Efficient Exact K-Means

arXiv: https://arxiv.org/abs/2603.09229 · HTML: https://mlsanity.com/paper/2603.09229/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 10 Mar 2026.

## Built on by (most cited first, 7 of 7)

- [SAW-INT4: System-Aware 4-Bit KV-Cache Quantization for Real-World LLM Serving](https://mlsanity.com/paper/2604.19157/) (2026)
- [CAKE: Compiler-Agent Co-Design for Frontier Kernel Evolution](https://mlsanity.com/paper/2608.12629/) (2026)
- [A Super Fast K-means for Indexing Vector Embeddings](https://mlsanity.com/paper/2603.20009/) (2026)
- [Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering](https://mlsanity.com/paper/2606.10896/) (2026)
- [PQ-HSA: Reusing Product-Quantized Scores for Hybrid Sparse-Approximate Attention](https://mlsanity.com/paper/2609.33746/) (2026)
- [GEM-KMeans: Memory-Efficient and Accurate Clustering on Massive Scale with GPU Optimization](https://mlsanity.com/paper/2609.36074/) (2026)
- [iCATS: Fast Video Generation via Interaction-Aware Sparse Attention and Timestep-Adaptive Sparsity](https://mlsanity.com/paper/2610.11302/) (2026)

## Abstract (arXiv)

k-means has historically been positioned primarily as an offline processing primitive, typically used for dataset organization or embedding preprocessing rather than as a first-class component in online systems. In this work, we revisit this classical algorithm under the lens of modern AI system design and enable k-means as an online primitive. We point out that existing GPU implementations of k-means remain fundamentally bottlenecked by low-level system constraints rather than theoretical algorithmic complexity. Specifically, the assignment stage suffers from a severe IO bottleneck due to the massive explicit materialization of the N × K distance matrix in High Bandwidth Memory (HBM). Simultaneously, the centroid update stage is heavily penalized by hardware-level atomic write contention caused by irregular, scatter-style token aggregations. To bridge this performance gap, we propose flash-kmeans, an IO-aware and contention-free k-means implementation for modern GPU workloads. Flash-kmeans introduces two core kernel-level innovations: (1) FlashAssign, which fuses distance computation with an online argmin to completely bypass intermediate memory materialization; (2) sort-inverse update, which explicitly constructs an inverse mapping to transform high-contention atomic scatters into high-bandwidth, segment-level localized reductions. Furthermore, we integrate algorithm-system co-designs, including chunked-stream overlap and cache-aware compile heuristics, to ensure practical deployability. Extensive evaluations on NVIDIA H200 GPUs demonstrate that flash-kmeans achieves up to 17.9× end-to-end speedup over best baselines, while outperforming industry-standard libraries like cuML and FAISS by 33× and over 200×, respectively. Our code is open-sourced at https://github.com/svg-project/flash-kmeans.