K-Search traduit CUDA pour Apple Silicon et frôle MLX

Berkeley et K-Search adaptent l’expertise CUDA aux puces Apple Silicon avec une extension du framework MLX. Les résultats publiés le 29 juillet 2026 montrent une vitesse atteignant 0,97 fois celle du noyau Attention natif de MLX. Sur le noyau Mamba SSM, l’implémentation atteint jusqu’à 20 fois le préremplissage de la version communautaire mlx-lm. K-Search ne recopie pas les instructions CUDA. Il traduit plutôt des stratégies d’optimisation vers les contraintes d’une autre architecture. Cette approche utilise des noyaux CUDA existants comme base de connaissances. Elle vise donc à réduire la reconstruction manuelle des optimisations sur chaque famille de puces. Lire l’étude complète.

Les noyaux concentrent une partie déterminante des performances des modèles. Ces programmes de bas niveau exécutent les opérations directement sur le processeur graphique. Les écrire efficacement demande généralement plusieurs années d’expertise. L’écosystème CUDA possède des implémentations optimisées de l’attention, des modèles à espace d’états et d’autres opérations centrales. MLX, le framework d’apprentissage automatique d’Apple, a progressé depuis la fin de 2023. Il exploite la mémoire unifiée des puces Apple Silicon pour l’inférence locale. Cette architecture convient notamment aux modèles de 7 à 70 milliards de paramètres sur les puces de la série M. Certaines opérations restent toutefois absentes ou peu optimisées.

K-Search explore les noyaux MLX

Le moteur évolutionnaire alterne raisonnement, génération de code, compilation et mesure sur le matériel réel. Un modèle de langage choisit d’abord une optimisation à tester. Un autre rôle produit ensuite une implémentation candidate. Dans les essais décrits, Gemini 3.5 Pro Preview assure ces deux fonctions. Le système classe le noyau, réécrit son calcul et identifie le goulot d’étranglement probable. Il peut s’agir de la bande passante, de la latence, du calcul ou de la synchronisation. K-Search conserve ces décisions dans un arbre de recherche. Chaque nœud reçoit une note globale entre 0 et 10, une confiance entre 0 et 1 et des impacts matériels estimés.

La couche de traduction relie les primitives CUDA à leurs équivalents Metal et MLX. La mémoire partagée CUDA devient une mémoire de groupe de travail Metal limitée à 32 Ko. Le débit HBM3 d’une H100, autour de 3,35 To par seconde, est comparé aux environ 400 Go par seconde de la mémoire unifiée d’une M3 Max. Cette différence modifie les optimisations prioritaires. La couche fournit aussi des motifs propres à Apple, comme les réductions par registres avec simd_shuffle_xor. Elle propose également l’astuce exp2, qui remplace exp(x) par exp2(x log2 e). Le noyau peut alors utiliser directement l’instruction matérielle rapide correspondante.

Mamba accélère le préremplissage

Le noyau Attention a été testé selon trois configurations sur Apple Silicon. La première utilise une version naïve. La deuxième applique l’évolution sans contexte supplémentaire. La troisième ajoute la couche de traduction complète. Avec ce contexte, le système retrouve plusieurs stratégies de FlashAttention-2. Il combine le découpage en mémoire de groupe, le softmax en ligne, la transposition de K et le double tampon. La performance progresse de 0,26 fois à 0,97 fois celle du noyau Attention natif de MLX. Les auteurs attribuent cette progression à la capacité de la couche à transmettre des stratégies d’implémentation. Le code généré reste néanmoins adapté aux primitives et limites d’Apple Silicon.

Mamba confirme l’extension à un calcul différent de l’attention. Son noyau SSM met à jour un état récurrent, sans dépendre d’un softmax. Sur un M1 Max de 64 Go, l’implémentation mlx-mamba atteint 152 jetons par seconde en décodage. En préremplissage, elle atteint 5 751 jetons par seconde avec une longueur de 512. Les mesures montent à 6 010, 6 612 et 6 743 jetons par seconde pour des longueurs de 1 024, 2 048 et 4 096. La version mlx-lm communautaire reste entre 326 et 339 jetons par seconde. La référence PyTorch mamba.py atteint entre 1 042 et 1 127 jetons par seconde.