Flash-KMeans: K-Means Esatto Veloce ed Efficiente in Termini di Memoria

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

March 10, 2026

Autori: 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

Abstract

Storicamente, il k-means è stato considerato principalmente una primitiva per l'elaborazione offline, tipicamente utilizzato per l'organizzazione di dataset o la pre-elaborazione di embedding, piuttosto che come un componente di prima classe nei sistemi online. In questo lavoro, rivalutiamo questo algoritmo classico attraverso la lente del moderno design dei sistemi di IA e abilitiamo il k-means come primitiva online. Sottolineiamo che le implementazioni esistenti del k-means su GPU rimangono fondamentalmente limitate da vincoli sistemistici di basso livello piuttosto che dalla complessità algoritmica teorica. Nello specifico, la fase di assegnazione soffre di un severo collo di bottiglia di I/O a causa della massiccia materializzazione esplicita della matrice delle distanze N per K nella High Bandwidth Memory (HBM). Simultaneamente, la fase di aggiornamento dei centroidi è fortemente penalizzata dalla contenzione hardware sulle scritture atomiche causata da aggregazioni irregolari di token di tipo scatter. Per colmare questo divario prestazionale, proponiamo flash-kmeans, un'implementazione del k-means consapevole dell'I/O e priva di contenzione per carichi di lavoro moderni su GPU. Flash-kmeans introduce due innovazioni fondamentali a livello di kernel: (1) FlashAssign, che fonde il calcolo della distanza con un argmin online per bypassare completamente la materializzazione intermedia in memoria; (2) l'aggiornamento "sort-inverse", che costruisce esplicitamente un mapping inverso per trasformare scatter atomici ad alta contenzione in riduzioni localizzate a livello di segmento ad alta larghezza di banda. Inoltre, integriamo co-progettazioni algoritmo-sistema, inclusa la sovrapposizione di stream in chunk ed euristiche di compilazione consapevoli della cache, per garantire la praticità del deployment. Valutazioni estensive su GPU NVIDIA H200 dimostrano che flash-kmeans raggiunge un miglioramento prestazionale end-to-end fino a 17.9 volte rispetto ai migliori baseline, superando rispettivamente di 33 volte e oltre 200 volte librerie standard del settore come cuML e FAISS.

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.