K-Search atteint 97 % de la vitesse native de MLX

Berkeley adapte K-Search pour Apple Silicon avec un backend MLX capable de traduire des kernels CUDA existants. Le framework utilise ces kernels comme base de connaissances, au lieu de reconstruire chaque optimisation. Les résultats publiés le 29 juillet 2026 montrent un kernel Attention atteignant 0,97 fois la vitesse du kernel natif de MLX. Sur le kernel SSM de Mamba, le préremplissage atteint jusqu’à 20 fois la vitesse de l’implémentation communautaire mlx-lm. Cette approche vise les opérations bas niveau qui déterminent les performances des modèles. L’équipe présente le travail dans le billet original de Berkeley.

MLX cible les Mac depuis la fin de 2023, selon les auteurs. Le framework exploite la mémoire unifiée des puces Apple Silicon pour l’inférence locale. Les auteurs indiquent que cette architecture convient notamment aux modèles de 7 à 70 milliards de paramètres sur les puces M. Ils signalent toutefois l’absence, ou la faible optimisation, de plusieurs kernels importants. La liste comprend notamment l’attention paginée, les scans SSM et le routage MoE fusionné. K-Search cherche donc à transférer l’expertise accumulée dans l’écosystème CUDA. Cette expertise regroupe des implémentations optimisées d’Attention et de modèles à espace d’états.

K-Search explore les kernels

Le framework combine recherche et code dans une boucle itérative. Un modèle de langage sélectionne une optimisation, puis un modèle générateur produit un kernel candidat. Le système compile ensuite ce candidat et le mesure sur le matériel réel. Les résultats alimentent la recherche suivante. Une spécification décrit les règles matérielles, les contraintes mathématiques et les motifs d’optimisation. Dans les expérimentations, Gemini 3.5 Pro Preview assure le raisonnement et l’écriture des kernels. Son modèle de décision conserve un arbre d’options. Chaque nœud reçoit une note globale de zéro à dix et un niveau de confiance compris entre zéro et un.

La traduction impose des contraintes propres à Apple Silicon. La mémoire threadgroup de Metal est limitée à 32 Ko, contre 48 Ko pour la référence NVIDIA citée. La bande passante annoncée atteint environ 400 Go/s sur un M3 Max, contre environ 3,35 To/s sur la mémoire HBM3 d’un H100. Ces écarts modifient les optimisations pertinentes. La couche de traduction associe des primitives CUDA à leurs équivalents MLX ou Metal. Elle fournit aussi des motifs spécifiques, comme les réductions par registres avec simd_shuffle_xor. Elle recommande également l’instruction exp2 pour accélérer le calcul du softmax.

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

L’Attention progresse par étapes selon la comparaison menée sur Apple Silicon. Le kernel naïf atteint 0,26 fois la vitesse du kernel Attention natif de MLX. L’évolution sans contexte supplémentaire sert de second point de comparaison. Avec la couche complète, le système découvre notamment le double tamponnage et le déroulage de boucle. Il retrouve aussi le tuilage en mémoire threadgroup et le softmax en ligne associés à FlashAttention-2. La transposition de K améliore les accès mémoire dans cette configuration. Les auteurs présentent ces résultats comme une mesure isolant l’effet de la couche de traduction, plutôt qu’une simple conversion syntaxique d’un kernel CUDA.