De CUDA à MLX : L'expertise de K-Search au service d'Apple Silicon

Une nouvelle ère dans le calcul

Le monde du calcul évolue à un rythme effréné. Les innovations matérielles ne se limitent pas à des GPU plus rapides, mais engendrent également une variété de puces de différents fabricants, chacune ayant sa propre architecture, souvent conçue pour des charges de travail spécifiques liées à l'intelligence artificielle (IA). En parallèle, le développement logiciel progresse tout aussi rapidement, avec des outils de codage en IA capables de générer ce qui prenait auparavant des mois en quelques minutes.

Dans un environnement où l'IA prend une place centrale, les kernels GPU deviennent essentiels. Ces programmes de bas niveau, exécutés au sein des GPU, requièrent une expertise considérable pour être optimisés efficacement. La migration d'un kernel d'un matériel à un autre représente un défi encore plus grand, nécessitant souvent une redécouverte des optimisations spécifiques à chaque architecture.

L'expertise des kernels CUDA

L'écosystème CUDA, soutenu par des décennies de recherche et d'optimisations, a permis de développer des implémentations hautement efficaces de fonctions critiques telles que l'attention ou les modèles d'espace d'état. Ces réalisations représentent des milliers d'heures d'ingénierie. Cependant, les nouvelles écosystèmes matériels, comme Apple Silicon et d'autres accélérateurs IA personnalisés, se développent rapidement mais manquent de cette profondeur d'expertise.

Il devient alors crucial de se demander si cette expertise peut être transférée de manière automatique.

K-Search : une avancée dans l'optimisation des kernels

K-Search, un cadre d'optimisation des kernels évolutif développé à l'UC Berkeley, se présente comme une réponse à ce défi. Cette méthode utilise l'intelligence artificielle pour optimiser des kernels GPU. En étendant K-Search avec un backend pour MLX, le framework d'apprentissage machine d'Apple, nous avons créé une couche de traduction CUDA vers MLX. Cette approche permet à K-Search de s'appuyer sur des kernels CUDA existants pour en générer de nouveaux pour Apple Silicon, réduisant ainsi le besoin de reconstruire des kernels à partir de zéro.

Les résultats parlent d'eux-mêmes : notre méthode atteint une performance proche de celle d'un expert sur Apple Silicon, offrant un gain de vitesse de 0,97x par rapport au kernel d'attention natif MLX, et jusqu'à 20x en vitesse de pré-remplissage par rapport à l'implémentation communautaire mlx-lm sur le kernel Mamba SSM.

Pourquoi choisir MLX ?

Depuis fin 2023, le framework MLX d'Apple a connu une adoption remarquable. Avec des millions de MacBooks et Mac Studios équipés d'Apple Silicon, MLX permet l'inférence IA locale sans coûts liés au cloud. Son architecture de mémoire unifiée le rend particulièrement adapté aux modèles de taille intermédiaire (de 7 à 70 milliards de paramètres sur les puces M).

Cependant, une lacune importante persiste : de nombreux kernels critiques pour la performance, souvent disponibles dans l'écosystème NVIDIA, se révèlent soit absents, soit peu optimisés sur MLX sans un réglage spécifique au matériel. Bien que MLX exécute correctement les modèles, il laisse souvent de côté un potentiel de performance considérable.

Le fonctionnement de K-Search

K-Search se distingue par sa capacité à itérer sur un kernel naïf en fonction d'une spécification matérielle. Un modèle d'IA raisonne sur les optimisations à appliquer, génère des kernels candidats, qui sont ensuite compilés et testés sur du matériel réel. Les mesures obtenues nourrissent le processus de recherche, permettant une amélioration continue jusqu'à ce que la performance converge.

Ce processus est guidé par un document spécifique à un domaine qui encode les règles matérielles et les contraintes mathématiques, empêchant la génération de code invalide et assurant que les candidats compilent et s'exécutent efficacement. En pratique, un modèle unique (comme Gemini 3.5 Pro Preview) joue à la fois le rôle de raisonnement et de génération de code, en analysant le kernel avant de proposer des optimisations.

Conclusion

En résumé, la combinaison de l'expertise accumulée dans l'écosystème CUDA et l'innovation apportée par K-Search ouvre de nouvelles perspectives pour l'optimisation des performances sur Apple Silicon. Les applications de cette méthode ne se limitent pas à MLX, mais s'étendent à tout écosystème où l'expertise CUDA peut être utilement transférée.

Si vous souhaitez en savoir plus sur l'optimisation des performances de votre infrastructure IA ou sur la mise en œuvre de stratégies similaires, n'hésitez pas à Contactez-moi.