Berkeley rapproche CUDA des performances natives de MLX

Berkeley adapte CUDA K-Search transpose une expertise accumulée dans l’écosystème CUDA vers les puces Apple Silicon. L’équipe a étendu ce cadre de recherche évolutionnaire avec un backend pour MLX, le framework d’apprentissage automatique d’Apple. Le système utilise des noyaux CUDA existants comme base de connaissances. Il les adapte ensuite en noyaux adaptés à l’architecture Apple, au lieu de repartir d’une implémentation vide. Les résultats atteignent 0,97 fois la vitesse du noyau Attention natif de MLX. Sur le noyau SSM de Mamba, le préremplissage progresse jusqu’à 20 fois face à l’implémentation communautaire mlx-lm. Lire l’article de Berkeley

MLX comble un écart MLX fonctionne sur les puces Apple Silicon présentes dans des centaines de millions de MacBook et Mac Studio. Son architecture à mémoire unifiée convient notamment aux modèles de 7 à 70 milliards de paramètres sur les puces de série M. Cependant, plusieurs noyaux importants restent absents ou peu optimisés. Le texte cite notamment l’attention paginée, les noyaux de balayage SSM et le routage MoE fusionné. K-Search cherche à réutiliser les optimisations développées pour CUDA. Cette approche vise donc une adaptation architecturale, et non une copie instruction par instruction des noyaux NVIDIA.

K-Search explore les noyaux

Une recherche guidée K-Search combine un modèle de raisonnement, un modèle de génération de code et des mesures réalisées sur le matériel réel. Dans les essais décrits, Gemini 3.5 Pro Preview assure les deux rôles. Le système classe d’abord le noyau, réécrit le calcul de référence et identifie le goulot d’étranglement probable. Il propose ensuite une seule modification par itération. Chaque candidat est compilé puis mesuré. Les résultats alimentent un arbre de recherche persistant. Chaque nœud reçoit une note globale de zéro à dix, un niveau de confiance et des impacts estimés sur la bande passante, les registres et l’adéquation matérielle.

La traduction impose des contraintes La couche CUDA vers MLX associe les primitives et leurs équivalents Metal avec des règles matérielles précises. La mémoire partagée CUDA devient la mémoire threadgroup Metal, limitée à 32 kilo-octets contre 48 kilo-octets pour NVIDIA. La bande passante H100 d’environ 3,35 téraocta-octets par seconde est comparée aux quelque 400 gigaoctets par seconde du M3 Max. La couche ajoute aussi des motifs spécifiques à MLX. Elle propose notamment des réductions par registres avec simd_shuffle_xor et l’utilisation de l’instruction exp2 pour accélérer le softmax. Voir le site de Shiyi Cao

Mamba accélère le préremplissage

Attention approche l’expertise Trois configurations ont été comparées sur un noyau Attention MLX. La première utilise une base naïve. La deuxième applique l’évolution sans contexte supplémentaire. La troisième fournit la couche de traduction complète. Cette dernière retrouve le tuilage en mémoire threadgroup, le softmax en ligne, la transposition de K et la technique exp2 associée à FlashAttention-2. La vitesse passe de 0,26 fois celle du noyau Attention Apple à 0,97 fois. Ce résultat mesure l’effet du contexte de traduction dans l’expérience présentée. Il ne correspond pas à une copie directe du noyau CUDA d’origine.

Mamba mesure le transfert L’équipe a ensuite testé le noyau SSM de Mamba sur un M1 Max de 64 gigaoctets, avec mamba-370m en précision f16. Le noyau mlx-mamba atteint 5 751 jetons par seconde en préremplissage à une longueur de 512. À 4 096, il atteint 6 743 jetons par seconde. L’implémentation mlx-lm atteint respectivement 329 et 339 jetons par seconde. En décodage, les mesures atteignent 152 jetons par seconde pour mlx-mamba, contre 116 pour mlx-lm et 40 pour mamba.py. Ces chiffres montrent une différence selon le type de calcul évalué.