Flash-KMeans : K-Means Exact Rapide et Économe en Mémoire

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

March 10, 2026

Auteurs: Shuo Yang, Haocheng Xi, Yilong Zhao, Muyang Li, Xiaoze Fan, Jintao Zhang, Han Cai, Yujun Lin, Xiuyu Li, Kurt Keutzer, Song Han, Chenfeng Xu, Ion Stoica

cs.AI

Résumé

Traditionnellement, k-means a été considéré principalement comme une primitive de traitement hors ligne, généralement utilisé pour l'organisation de jeux de données ou le prétraitement d'embeddings, plutôt que comme un composant à part entière dans les systèmes en ligne. Dans ce travail, nous revisitons cet algorithme classique sous l'angle de la conception moderne des systèmes d'IA et permettons à k-means de fonctionner comme une primitive en ligne. Nous soulignons que les implémentations existantes de k-means sur GPU restent fondamentalement limitées par des contraintes système de bas niveau plutôt que par la complexité algorithmique théorique. Plus précisément, l'étape d'affectation souffre d'un goulot d'étranglement sévère des E/S dû à la matérialisation explicite massive de la matrice de distance N fois K dans la mémoire à haute bande passante (HBM). Simultanément, l'étape de mise à jour des centroïdes est fortement pénalisée par les conflits d'écriture atomique au niveau matériel, causés par des agrégations de tokens irrégulières de type scatter. Pour combler cet écart de performance, nous proposons flash-kmeans, une implémentation de k-means consciente des E/S et sans contention pour les charges de travail modernes sur GPU. Flash-kmeans introduit deux innovations fondamentales au niveau du noyau : (1) FlashAssign, qui fusionne le calcul de distance avec un argmin en ligne pour contourner complètement la matérialisation intermédiaire en mémoire ; (2) la mise à jour par inverse-tri, qui construit explicitement un mapping inverse pour transformer des scatters atomiques à haute contention en réductions localisées au niveau des segments, à haute bande passante. De plus, nous intégrons des co-conceptions algorithme-système, incluant le chevauchement par flux en tronçons et des heuristiques de compilation sensibles au cache, pour garantir une déployabilité pratique. Des évaluations approfondies sur les GPU NVIDIA H200 démontrent que flash-kmeans atteint une accélération de bout en bout allant jusqu'à 17,9 fois par rapport aux meilleures bases de référence, tout en surpassant les bibliothèques standard de l'industrie comme cuML et FAISS par des facteurs de 33 fois et plus de 200 fois, respectivement.

English

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 times 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.9times end-to-end speedup over best baselines, while outperforming industry-standard libraries like cuML and FAISS by 33times and over 200times, respectively.