2MarkTechPost
Flash-KMeans : un K-Means exact et optimisé pour les E/S, plus de 200 fois plus rapide que FAISS sur GPU
Des chercheurs de l'Université de Californie à Berkeley et de l'Université du Texas à Austin ont publié Flash-KMeans, une bibliothèque open source qui réimplémente l'algorithme k-means standard de Lloyd sur GPU, avec des gains de performance spectaculaires. Sur un NVIDIA H200, la bibliothèque affiche jusqu'à 17,9 fois plus de rapidité que le meilleur concurrent testé, 33 fois plus que la bibliothèque industrielle cuML de NVIDIA, et plus de 200 fois plus que FAISS, la référence du secteur pour la recherche vectorielle. Flash-KMeans s'installe via pip et est distribué sous licence Apache 2.0. Le résultat mathématique est strictement identique à un k-means classique : aucune approximation, aucun raccourci algorithmique. L'enjeu est de taille parce que le k-means n'est plus seulement un outil de prétraitement utilisé une fois avant l'entraînement. Les pipelines d'IA modernes l'appellent en boucle, à l'intérieur même des phases d'entraînement et d'inférence, ce qui rend chaque milliseconde critique. Flash-KMeans attaque deux goulots d'étranglement distincts. La phase d'assignation, qui consiste à associer chaque point au centroïde le plus proche, génère habituellement une matrice de distances de taille N x K entièrement écrite en mémoire HBM avant d'être relue : sur N=65 536 points, K=1 024 clusters et d=128 dimensions, le calcul arithmétique prend 2,6 ms mais les accès mémoire coûtent 23 ms. La solution, baptisée FlashAssign, s'inspire de FlashAttention : elle fusionne le calcul de distance et la recherche du minimum en tuiles traitées sur la SRAM on-chip, sans jamais matérialiser la matrice complète. La phase de mise à jour des centroïdes, elle, souffrait de collisions atomiques massives sur les clusters populaires, limitant la bande passante effective à 50 Go/s sur le H200. La méthode Sort-Inverse Update contourne ce problème en triant les assignations par identifiant de cluster, ce qui permet de réduire chaque segment localement avant une seule opération atomique par cluster. Flash-KMeans s'inscrit dans une dynamique plus large où les optimisations de bas niveau, au niveau du noyau GPU, deviennent aussi décisives que les innovations algorithmiques. La bibliothèque FAISS, développée par Meta et omniprésente dans les systèmes de recherche vectorielle en production, ne passe pas à l'échelle sans compromis : les implémentations PyTorch classiques tombent en panne mémoire dès que K devient grand, faute de pouvoir matérialiser la matrice N x K. Flash-KMeans traite un milliard de points avec K=32 768 et d=128 en 41,4 secondes contre 261,8 secondes pour la référence, et ce hors-coeur. Avec la montée en puissance des bases de données vectorielles et du clustering dynamique dans les systèmes RAG et de recommandation, une implémentation exacte et aussi rapide pourrait rapidement devenir un composant standard des pipelines d'IA à grande échelle.
UELes laboratoires et entreprises européens déployant des systèmes RAG ou des bases de données vectorielles à grande échelle peuvent bénéficier directement de cette bibliothèque open source pour accélérer leurs pipelines de clustering sans modification algorithmique.
