Decades of CUDA optimization transferred automatically to Apple silicon

A research team extended K-Search, an evolutionary framework for kernel optimization originally developed at UC Berkeley Sky Lab, so that it can ingest existing CUDA kernels and translate them to MLX, Apple’s machine-learning framework for its silicon chips. Rather than rediscovering optimizations from scratch for each hardware platform, the method treats the existing CUDA code as a knowledge base and adapts it to the new architecture.

K-Search receives a naïve kernel and a hardware specification, then runs an iterative optimization loop. A language model evaluates which optimizations are worth trying, a code-generation model produces candidate kernels, and those candidates are compiled and benchmarked on real hardware. The measurements are fed back into the search, which refines promising directions and discards dead ends until performance converges. The search is guided by a Spec document that encodes hardware rules, optimization patterns, and mathematical constraints, preventing the generated code from violating illegal primitives and ensuring that candidates compile and run efficiently. In the experiments, a single model—Gemini 3.5 Pro Preview—performed both analysis and code generation.

The results show that the method attains near-expert performance on Apple silicon, achieving a speed-up of 0.97× compared with the built-in MLX Attention kernel and up to 20× on the Mamba SSM kernel relative to the community implementation mlx-lm. The study quantifies how much of the gain stems from the translation layer itself versus the search loop. Although the work focuses on MLX, the approach is not specific to MLX and applies to any environment where CUDA knowledge can be transferred.

MLX has seen broad adoption since late 2023. With Apple silicon deployed in hundreds of millions of MacBook and Mac Studio devices, MLX enables local model execution without cloud costs, and its unified memory architecture makes it attractive for medium-sized models (7B to 70B parameters on the M series chips). Beneath this momentum lies a significant gap: critical performance kernels that NVIDIA treats as built-in—such as paged attention, SSM-optimized scan kernels, or fused MoE routing—are either absent from MLX or exist only as naïve implementations lacking hardware-specific tuning. While MLX can run models correctly, it leaves considerable performance on the table. That gap motivates the present work.