Flash-KMeans réduit la latence du k-means sur GPU grâce à une gestion optimisée des donnée
Les pipelines d’IA modernes intègrent désormais des appels au k-means dans leurs boucles d’entraînement et d’inférence, où chaque milliseconde compte. Des chercheurs de l’UC Berkeley et de l’UT Austin viennent de publier en open source Flash-KMeans, une implémentation adaptée aux entrées/sorties qui préserve les calculs de l’algorithme de Lloyd tout en réorganisant le flux de données sur GPU NVIDIA. Résultat : jusqu’à 17,9× de gain en temps d’exécution par rapport aux meilleures références sur un H200, 33× face à NVIDIA cuML, et plus de 200× comparé à FAISS.
Noyaux optimisés, sans compromis
Contrairement aux méthodes qui réduisent le travail par élagage ou échantillonnage, Flash-KMeans conserve tous les centroïdes et points dans les calculs. Il fusionne simplement la boucle de distance avec un argmin en ligne, évitant ainsi de transférer la matrice complète de distances N×K depuis la mémoire haute bande passante. S’inspirant des principes de FlashAttention, le noyau FlashAssign traite des blocs de points et centroïdes depuis la HBM vers la SRAM, réduisant la complexité d’E/S de O(NK) à O(Nd + Kd). L’équipe de recherche rapporte des gains jusqu’à 21,2× au niveau des noyaux, ramenant par exemple un temps d’affectation de 122,5 ms à seulement 5,8 ms lors d’un test.
Mises à jour sans contention
Dans les implémentations classiques, la phase de mise à jour des centroïdes repose sur des opérations atomiques de type diffusion qui sérialisent lorsque de nombreux threads ciblent le même cluster « chaud ». Flash-KMeans remplace cette approche par Sort-Inverse Update : il trie le vecteur d’affectation 1D pour regrouper les identifiants de clusters identiques en segments contigus, puis chaque bloc de threads réduit un segment en SRAM avant d’effectuer une seule opération atomique. Cette méthode diminue le nombre d’opérations atomiques et augmente la bande passante effective, passant d’environ 50 Go/s à des niveaux bien supérieurs, offrant jusqu’à 6,3× d’accélération au niveau des noyaux.
Installable via pip sous licence Apache 2.0, Flash-KMeans cible les charges de travail en k-means par lots où la latence par appel prime sur le pic de FLOPS.
Source : MarkTechPost. Synthèse éditoriale assistée par IA — TechnoExpress.