Flash-KMeans: Schneller und speichereffizienter exakter K-Means-Algorithmus
Flash-KMeans: Fast and Memory-Efficient Exact K-Means
March 10, 2026
Autoren: 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
Zusammenfassung
K-Means wurde historisch primär als Offline-Verarbeitungsprimitive positioniert, die typischerweise zur Datensatzorganisation oder Embedding-Vorverarbeitung eingesetzt wird, anstatt als First-Class-Komponente in Online-Systemen. In dieser Arbeit betrachten wir diesen klassischen Algorithmus neu durch die Linse moderner KI-Systemgestaltung und ermöglichen K-Means als Online-Primitive. Wir weisen darauf hin, dass bestehende GPU-Implementierungen von K-Means nach wie vor grundlegend durch Low-Level-Systembeschränkungen und nicht durch theoretische algorithmische Komplexität limitiert werden. Insbesondere leidet der Zuordnungsschritt (Assignment Stage) unter einem schwerwiegenden E/A-Engpass (I/O-Bottleneck), der auf die massive explizite Materialisierung der N-mal-K-Distanzmatrix im High-Bandwidth-Memory (HBM) zurückzuführen ist. Gleichzeitig wird der Zentroiden-Aktualisierungsschritt (Centroid Update Stage) erheblich durch hardwarebedingte Schreibkonflikte (Atomic Write Contention) beeinträchtigt, die durch unregelmäßige, streuende (Scatter-style) Token-Aggregationen verursacht werden. Um diese Leistungslücke zu schließen, schlagen wir flash-kmeans vor, eine E/A-bewusste (IO-aware) und kontentionsfreie K-Means-Implementierung für moderne GPU-Workloads. Flash-kmeans führt zwei zentrale Kernel-Level-Innovationen ein: (1) FlashAssign, das die Distanzberechnung mit einer Online-Argmin-Operation fusioniert, um die Zwischenspeicher-Materialisierung vollständig zu umgehen; (2) Sort-Inverse-Update, das explizit eine inverse Abbildung konstruiert, um hochkontentionierte atomare Streuoperationen (Atomic Scatters) in hochbandbreitige, segmentlokale Reduktionen (Segment-level Localized Reductions) umzuwandeln. Darüber hinaus integrieren wir Algorithmus-System-Co-Designs, einschließlich chunked-stream overlap und cache-aware compile heuristics, um die praktische Einsatzfähigkeit sicherzustellen. Umfangreiche Evaluierungen auf NVIDIA H200 GPUs zeigen, dass flash-kmeans eine bis zu 17,9-fache Ende-zu-Ende-Beschleunigung gegenüber den besten Baseline-Verfahren erreicht und dabei branchenübliche Bibliotheken wie cuML und FAISS um das 33-fache bzw. über 200-fache übertrifft.
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.