Flash-KMeans: K-Means Exato Rápido e Eficiente em Memória
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
Resumo
O k-means tem sido historicamente posicionado principalmente como um primitivo de processamento offline, tipicamente utilizado para organização de conjuntos de dados ou pré-processamento de *embeddings*, em vez de um componente de primeira classe em sistemas online. Neste trabalho, revisitamos este algoritmo clássico sob a ótica do projeto moderno de sistemas de IA e capacitamos o k-means como um primitivo online. Salientamos que as implementações existentes de k-means para GPU permanecem fundamentalmente limitadas por restrições de baixo nível do sistema, e não pela complexidade algorítmica teórica. Especificamente, a etapa de atribuição sofre de um severo gargalo de E/S devido à massiva materialização explícita da matriz de distâncias N vezes K na Memória de Alta Largura de Banda (HBM). Simultaneamente, a etapa de atualização de centróides é fortemente penalizada pela contenção de escritas atômicas a nível de *hardware*, causada por agregações de *tokens* irregulares, do tipo *scatter*. Para superar esta lacuna de desempenho, propomos o *flash-kmeans*, uma implementação de k-means consciente da E/S e livre de contenção para cargas de trabalho modernas em GPU. O *flash-kmeans* introduz duas inovações centrais a nível de *kernel*: (1) *FlashAssign*, que funde o cálculo da distância com um *argmin* online para contornar completamente a materialização intermediária na memória; (2) atualização por inversão e ordenação (*sort-inverse update*), que constrói explicitamente um mapeamento inverso para transformar operações de *scatter* atômicas de alta contenção em reduções localizadas a nível de segmento, de alta largura de banda. Adicionalmente, integramos co-projetos de algoritmo e sistema, incluindo sobreposição de *stream* em blocos (*chunked-stream overlap*) e heurísticas de compilação conscientes da *cache*, para garantir a implantabilidade prática. Avaliações extensivas em GPUs NVIDIA H200 demonstram que o *flash-kmeans* até 17,9 vezes de aceleração de ponta a ponta em relação às melhores *baselines*, superando também bibliotecas padrão da indústria como cuML e FAISS em 33 vezes e mais de 200 vezes, respetivamente.
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.