Flash-KMeans: Быстрый и экономичный по памяти точный алгоритм K-средних

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

March 10, 2026

Авторы: 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

Аннотация

Исторически k-means позиционировался в первую очередь как примитив для офлайн-обработки, обычно используемый для организации наборов данных или предобработки эмбеддингов, а не как полноценный компонент онлайн-систем. В данной работе мы пересматриваем этот классический алгоритм через призму современного проектирования ИИ-систем и реализуем k-means как онлайн-примитив. Мы указываем, что существующие GPU-реализации k-means остаются фундаментально ограниченными низкоуровневыми системными ограничениями, а не теоретической алгоритмической сложностью. В частности, этап назначения страдает от серьезного узкого места ввода-вывода из-за массового явного материализации матрицы расстояний размером N на K в высокоскоростной памяти (HBM). Одновременно этап обновления центроидов сильно замедляется из-за аппаратной конкуренции атомарных операций записи, вызванной нерегулярными, scatter-подобными агрегациями токенов. Для преодоления этого разрыва в производительности мы предлагаем flash-kmeans, IO-оптимизированную и бесконфликтную реализацию k-means для современных GPU-нагрузок. Flash-kmeans представляет два ключевых нововведения на уровне ядер: (1) FlashAssign, который объединяет вычисление расстояний с онлайн-поиском argmin, чтобы полностью избежать промежуточной материализации в памяти; (2) sort-inverse update, который явно строит обратное отображение для преобразования высококонфликтных атомарных scatter-операций в высокоскоростные локализованные редукции на уровне сегментов. Кроме того, мы интегрируем совместные алгоритмико-системные решения, включая чанковое перекрытие потоков и кэш-ориентированные эвристики компиляции, чтобы обеспечить практическую применимость. Масштабные оценки на GPU NVIDIA H200 демонстрируют, что flash-kmeans обеспечивает до 17.9-кратного сквозного ускорения по сравнению с лучшими базовыми методами, превосходя отраслевые стандартные библиотеки, такие как cuML и FAISS, в 33 и более 200 раз соответственно.

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.