Flash-KMeans: algorytm k-means 200× szybszy niż FAISS na GPU

Zespół badaczy zaprezentował Flash-KMeans – nową implementację algorytmu k-means zoptymalizowaną pod kątem pracy na procesorach GPU. Kod jest dostępny otwarcie i napisany w kernelach Triton, przy czym algorytm nie przybliża matematyki ani nie zmienia samego działania – pozostaje dokładnie taki jak oryginalny Lloyd's k-means, tylko znacznie szybszy.

Szczególnie istotne są dwie innowacyjne techniki optymalizacji. FlashAssign eliminuje konieczność materializacji pełnej macierzy odległości, co oszczędza olbrzymie ilości pamięci. Z kolei Sort-Inverse Update pozbywa się problemu atomowego konfliktu dostępu do pamięci. Te rozwiązania skupiają się na tym, co w branży GPU zwie się IO-aware computing – czyli świadomym zarządzaniem przepływem danych między różnymi poziomami pamięci.

Wyniki benchmarków są imponujące: na procesorze NVIDIA H200 Flash-KMeans osiąga przyspieszenie 17,9× w stosunku do innych implementacji end-to-end, 33× szybciej niż cuML, a aż ponad 200× szybciej niż FAISS – de facto standardu branży. To ma realny wpływ na praktyczne zastosowania: analiza dużych zbiorów danych, wyszukiwanie podobieństwa i segmentacja mogą być teraz wykonane w ułamku poprzedniego czasu.