Une équipe de recherche vient de publier Flash-KMeans, une bibliothèque ouverte qui réalise le partitionnement K-means exact plus de 200 fois plus rapidement que FAISS sur GPU, sans rien changer aux mathématiques. Le K-means est l'un des algorithmes de partitionnement les plus anciens et les plus utilisés, le cheval de bataille derrière tout, de l'indexation des bases de données vectorielles à la compression d'images, et Flash-KMeans ne l'approxime pas et ne modifie pas la procédure classique de Lloyd. Il repense plutôt la façon dont l'algorithme déplace les données à travers un GPU moderne, exactement l'intuition qui a rendu FlashAttention si rapide: le goulot d'étranglement, ce n'est pas le calcul, c'est le trafic mémoire.
L'accélération repose sur deux idées au niveau des noyaux. La première, FlashAssign, fusionne le calcul des distances avec un argmin en ligne pour que l'étape d'assignation n'ait jamais à écrire une énorme matrice intermédiaire de distances en mémoire pour ensuite la relire; le travail reste dans la mémoire rapide intégrée à la puce. La seconde, appelée sort-inverse update, s'attaque à l'autre étape coûteuse, soit le recalcul du centre de chaque regroupement. Réalisée naïvement, cette étape met des milliers de fils du GPU en concurrence pour additionner dans les mêmes quelques accumulateurs, un schéma très conflictuel qui bride le matériel. Flash-KMeans construit plutôt une cartographie inverse explicite qui transforme ces écritures atomiques dispersées en réductions ordonnées, à haut débit, par segments.
Sur un NVIDIA H200, l'équipe rapporte jusqu'à 17,9 fois d'accélération de bout en bout par rapport à la meilleure référence antérieure, 33 fois plus rapidement que le propre cuML de NVIDIA, et plus de 200 fois plus rapidement que FAISS, la bibliothèque que la plupart des gens utilisent quand ils ont besoin de K-means à grande échelle. Et c'est exact: l'algorithme produit les mêmes regroupements que ceux de Lloyd, et non une approximation plus rapide qui sacrifierait discrètement la précision. La bibliothèque est écrite en Triton, le langage à la Python conçu pour les noyaux GPU, est distribuée sous la licence permissive Apache 2.0 et s'installe avec une seule commande pip. L'article est sur arXiv et le code, sur GitHub.
Le résultat est un petit exemple bien net d'un schéma qui continue de porter ses fruits en systèmes d'apprentissage automatique: les plus gros gains viennent souvent non pas d'un nouvel algorithme, mais du fait d'en réimplémenter un ancien pour respecter la hiérarchie mémoire du matériel sur lequel il tourne. FlashAttention l'a fait pour l'attention des transformeurs, et Flash-KMeans le fait pour le partitionnement, une primitive qui se cache sous la recherche d'information, la déduplication et la recherche vectorielle qui propulse une bonne partie des applications d'IA d'aujourd'hui. Une accélération de 200 fois sur un algorithme exact et inchangé, ce n'est pas une optimisation marginale, c'est la différence entre partitionner un jeu de données en quelques secondes et patienter de longues minutes. La réserve honnête, c'est que les multiplicateurs annoncés dépendent du matériel, du jeu de données et de la configuration de la référence, alors les gains réels varieront; mais la direction est claire, et le code est là, juste devant, pour qu'on le vérifie.