K-Search accélère Mamba jusqu’à 20 fois sur Apple Silicon

K-Search adapte l’expertise CUDA aux puces Apple Silicon via un backend MLX. L’équipe a ajouté cette capacité au framework de recherche évolutive présenté par Shiyi Cao au Sky Lab de Berkeley. Le système utilise des kernels CUDA existants comme base de connaissances. Il les transforme en stratégies adaptées à l’architecture Apple, au lieu de copier chaque instruction. Sur le kernel Attention, la configuration complète atteint 0,97 fois la vitesse du kernel Attention natif de MLX. Sur le kernel SSM de Mamba, elle obtient jusqu’à 20 fois le débit de préremplissage de l’implémentation communautaire mlx-lm. Les résultats sont détaillés dans l’article de Berkeley.

MLX équipe les puces Apple Silicon pour l’inférence locale. Le framework bénéficie d’une architecture mémoire unifiée et cible notamment les modèles de 7 à 70 milliards de paramètres sur les puces de la série M. Pourtant, plusieurs kernels spécialisés restent absents ou peu optimisés. Le texte cite notamment l’attention paginée, les kernels SSM et le routage MoE fusionné. CUDA dispose déjà de décennies d’optimisations manuelles pour l’attention et les modèles à espace d’états. Cette expertise représente des milliers d’heures d’ingénierie. K-Search cherche donc à transférer ces méthodes vers MLX, plutôt qu’à reconstruire chaque optimisation depuis une base naïve. Le projet s’appuie aussi sur le dépôt GitHub de Gal Bloch.

K-Search explore les kernels MLX

Le moteur évolutif alterne raisonnement, génération de code, compilation et mesure sur le matériel réel. Un modèle de langage propose une optimisation, puis un modèle de génération produit une variante exécutable. Les performances mesurées alimentent ensuite la recherche. Dans les essais décrits, Gemini 3.5 Pro Preview joue ces deux rôles. Le système classe d’abord le kernel, réécrit le calcul sous forme canonique et identifie le goulot d’étranglement probable. Il conserve ensuite un arbre de décisions appelé modèle du monde. Chaque nœud reçoit une note globale de zéro à dix, un niveau de confiance et des impacts estimés sur la mémoire, les registres et le calcul.

La traduction structurée relie les primitives CUDA aux équivalents Metal et MLX. La mémoire partagée CUDA devient la mémoire threadgroup de Metal, limitée à 32 kilo-octets contre 48 kilo-octets pour l’exemple NVIDIA cité. La réduction warp est associée à MMA, tandis que la synchronisation devient une barrière threadgroup. Le débit mémoire H100 indiqué atteint environ 3,35 téraocta-octets par seconde, contre environ 400 gigaoctets par seconde pour le M3 Max. Cette différence modifie les optimisations prioritaires. La couche fournit aussi des motifs spécifiques, comme les réductions par registres avec simd_shuffle_xor et l’utilisation de exp2 pour accélérer le softmax.

Mamba dépasse mlx-lm sur M1 Max

Le kernel Attention a été comparé selon trois configurations. La première utilise une base naïve. La deuxième applique uniquement l’évolution. La troisième ajoute la couche de traduction complète. Cette dernière retrouve plusieurs stratégies de FlashAttention 2, dont le tuilage en mémoire threadgroup, le softmax en ligne, la transposition de K et le double tampon. Elle atteint 0,97 fois la vitesse du kernel Attention natif de MLX, contre 0,26 fois pour le niveau initial présenté. L’optimisation exp2 remplace chaque exponentielle naturelle par une exponentielle en base deux, avec la relation exacte indiquée dans l’étude. Le kernel peut ainsi appeler l’instruction matérielle fast::exp2().