Flash-KMeans: точный k-means в 200x быстрее

Исследователи из Калифорнийского университета в Беркли и Техасского университета в Остине выпустили в открытый доступ Flash-KMeans — реализацию классической кластеризации k-means (алгоритм Ллойда) на GPU, которая, по их данным, работает более чем в 200 раз быстрее FAISS. Главное: она «не меняет математику и не использует приближений. Она лишь перестраивает то, как алгоритм перемещает данные на GPU». Лицензия Apache 2.0, установка одной командой 'pip install flash-kmeans'.

Замеры на NVIDIA H200 в FP16 при размерности 128 впечатляют: в 33 раза быстрее собственной cuML от NVIDIA, до 17,9x быстрее лучшего из аналогов в сквозном тесте и до 10,5x быстрее fastkmeans на «вне-памяти» прогоне по 1 миллиарду точек. Два собственных ядра дают свои ускорения — до 21,2x и 6,3x — которые и складываются в итоговые цифры.

За счёт чего ускорение

Работают два приёма с памятью. FlashAssign объединяет расчёт расстояний с онлайновым argmin, так что полная матрица расстояний N×K вообще не выгружается в память — сложность по вводу-выводу падает с O(NK) до O(Nd+Kd). Sort-Inverse Update сортирует точки по ID кластера, превращая разрозненные атомарные записи в непрерывные сегментные свёртки и снижая конкуренцию за атомарные операции. Ни один шаг не трогает результат: кластеры те же, просто данные гоняются по памяти GPU гораздо экономнее.

Почему это важно

K-means — это незаметная «сантехника» под многими современными системами ИИ: индексация для векторного поиска, маршрутизация в разреженном внимании, сжатие KV-кэша, низкобитная квантизация KV, диффузионные трансформеры. Такие конвейеры вызывают кластеризацию снова и снова внутри циклов обучения и инференса, где реально бьёт по карману задержка на вызов, а не теоретическая сложность. Прямая замена с тем же результатом, но быстрее, означает те же ответы за меньшее время и деньги — и спорить о потере точности не о чем, потому что метод точный.