Flash-KMeans: K-Medias Exacto Rápido y Eficiente en Memoria
Flash-KMeans: Fast and Memory-Efficient Exact K-Means
March 10, 2026
Autores: 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
Resumen
Históricamente, k-means se ha posicionado principalmente como una primitiva de procesamiento offline, utilizada típicamente para la organización de conjuntos de datos o el preprocesamiento de *embeddings*, en lugar de ser un componente de primera clase en sistemas online. En este trabajo, reexaminamos este algoritmo clásico bajo la lente del diseño moderno de sistemas de IA y habilitamos k-means como una primitiva online. Señalamos que las implementaciones existentes de k-means en GPU siguen estando fundamentalmente limitadas por restricciones de bajo nivel del sistema, más que por la complejidad algorítmica teórica. Específicamente, la etapa de asignación sufre un severo cuello de botella de E/S debido a la materialización explícita masiva de la matriz de distancias N por K en la Memoria de Ancho de Banda Alto (HBM). Simultáneamente, la etapa de actualización de centroides se ve fuertemente penalizada por la contención de escrituras atómicas a nivel de hardware, causada por agregaciones de tokens irregulares del tipo *scatter*. Para salvar esta brecha de rendimiento, proponemos flash-kmeans, una implementación de k-means consciente de la E/S y libre de contención para cargas de trabajo modernas en GPU. Flash-kmeans introduce dos innovaciones centrales a nivel de *kernel*: (1) FlashAssign, que fusiona el cálculo de distancias con un *argmin* online para evitar por completo la materialización intermedia en memoria; (2) la actualización *sort-inverse*, que construye explícitamente un mapeo inverso para transformar las operaciones de *scatter* atómicas de alta contención en reducciones localizadas a nivel de segmento y de alto ancho de banda. Además, integramos co-diseños algoritmo-sistema, incluyendo solapamiento de transmisión por fragmentos (*chunked-stream overlap*) y heurísticas de compilación conscientes de la caché, para garantizar la capacidad de implementación práctica. Evaluaciones exhaustivas en GPUs NVIDIA H200 demuestran que flash-kmeans logra una aceleración de extremo a extremo de hasta 17.9 veces sobre los mejores métodos de referencia, superando a bibliotecas estándar de la industria como cuML y FAISS en 33 veces y más de 200 veces, respectivamente.
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.