K-Search adapte CUDA et frôle l’Attention native de MLX

K-Search adapte des kernels CUDA à Apple Silicon grâce à une extension du framework MLX. L’équipe a construit une couche de traduction structurée pour réutiliser des connaissances d’optimisation existantes. Cette méthode atteint 0,97 fois la vitesse du kernel Attention natif de MLX. Elle obtient aussi jusqu’à 20 fois plus de vitesse en préremplissage sur le kernel Mamba SSM, comparé à l’implémentation communautaire mlx-lm. Le projet s’appuie sur K-Search, un framework évolutif développé à Berkeley Sky Lab. Celui-ci demande à des modèles d’IA de proposer, générer, compiler et mesurer des kernels sur le matériel réel. La présentation complète est disponible sur le blog de Berkeley.

MLX cible les puces Apple Silicon utilisées dans des centaines de millions de MacBook et Mac Studio. Son architecture mémoire unifiée facilite l’inférence locale de modèles comptant entre 7 et 70 milliards de paramètres sur les puces de la série M. Le framework exécute correctement ces modèles, mais certains kernels restent absents ou peu optimisés. Le billet cite notamment l’attention paginée, les scans SSM optimisés et le routage MoE fusionné. K-Search cherche donc à transférer des décennies d’expertise CUDA vers cette architecture. L’objectif n’est pas de recopier les instructions. La méthode traduit plutôt les stratégies vers les contraintes propres à Apple Silicon et à MLX.

K-Search traduit les stratégies CUDA

La recherche évolutive combine un modèle de raisonnement et un modèle de génération de code. Dans les expériences, Gemini 3.5 Pro Preview assure ces deux rôles. Le système classe d’abord le kernel, réécrit le calcul de référence et identifie le goulot d’étranglement probable. Il propose ensuite une seule optimisation par itération. Chaque hypothèse rejoint un arbre de décision persistant. Les nœuds reçoivent une note globale de zéro à dix, un niveau de confiance entre zéro et un, ainsi que des impacts estimés sur la bande passante et les registres. La recherche conserve les branches prometteuses et abandonne celles qui stagnent après plusieurs cycles.

La couche structurée associe des primitives CUDA à leurs équivalents Metal ou MLX. La mémoire __shared__ devient une mémoire de groupe limitée à 32 kilo-octets, contre 48 kilo-octets pour la référence NVIDIA citée. __syncthreads() devient threadgroup_barrier. La bande passante H100 d’environ 3,35 téraoctets par seconde est comparée aux quelque 400 gigaoctets par seconde du 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 dans une tuile MMA 8×8. Elle encode enfin des assertions réutilisables, plutôt que du code CUDA copié directement.

Le kernel Mamba accélère le préremplissage

Le kernel Attention a été évalué selon trois configurations sur Apple Silicon. La première utilise une base naïve. La deuxième applique l’évolution sans contexte supplémentaire. La troisième fournit la couche complète de traduction. Cette dernière redécouvre le tuilage en mémoire de groupe, le softmax en ligne, la transposition de K et l’astuce exp2. Le kernel évolué passe ainsi de 0,26 fois à 0,97 fois la vitesse du kernel Attention natif de MLX. L’astuce remplace chaque exponentielle par une exponentielle en base deux. Elle utilise directement l’instruction matérielle rapide fast::exp2(). Les détails techniques sont présentés dans le projet de Shiyi Cao.

Le test Mamba mesure un kernel SSM sur un M1 Max doté de 64 gigaoctets de mémoire. Avec un modèle mamba-370m en f16, l’implémentation mlx-mamba atteint 152 tokens par seconde en décodage. Elle atteint 5 751 tokens par seconde en préremplissage avec une longueur de 512. Les résultats montent à 6 010, 6 612 et 6 743 tokens par seconde pour des longueurs de 1 024, 2 048 et 4 096. L’implémentation mlx-lm reste entre 326 et 339 tokens par seconde sur ces quatre longueurs. La référence PyTorch mamba.py atteint entre 1 042 et 1 127 tokens par seconde.

La portée annoncée dépasse l’écosystème MLX, car la méthode vise les environnements capables de recevoir une expertise CUDA traduisible. Le backend développé pour Apple Silicon prend en charge la compilation et l’exécution via les interfaces Metal et C++ de MLX. Il intègre aussi les mesures de performance de mlx.core et des instructions adaptées à la génération de kernels Metal. Les travaux antérieurs de K-Search portaient sur quatre kernels CUDA de FlashInfer. Sur un budget de 120 itérations, le framework obtenait des améliorations plus régulières qu’OpenEvolve et ShinkaEvolve selon les résultats reproduits. L’approche conserve donc une boucle expérimentale fondée sur des mesures matérielles réelles.