From CUDA to MLX: The Bridge Bringing Decades of Optimization to Apple Silicon

What happened?

The Berkeley Sky Lab team, in collaboration with IBM Research, has extended their evolutionary kernel search framework K-Search with a structured translation layer from CUDA to MLX. This layer allows expert knowledge accumulated in CUDA kernels—such as FlashAttention or Mamba—to be automatically transferred to the Apple Silicon ecosystem, avoiding the need to rediscover each optimization manually. The results are compelling: they achieve 97% of the native MLX kernel performance for FlashAttention, and up to 20x speedup in the prefill of the Mamba state space model compared to the community implementation mlx-lm.

The work, published on the BAIR (Berkeley Artificial Intelligence Research) blog, details how the K-Search framework, originally focused on evolutionary search for CUDA kernels, has been adapted to generate optimized kernels for MLX, Apple's machine learning framework for its silicon. The key is a translation layer that maps CUDA optimizations to native MLX strategies, leveraging the unique features of Apple Silicon's unified memory architecture. According to the authors, this approach not only accelerates development but also allows the Apple ecosystem to benefit from decades of CUDA kernel engineering.

Why is it important?

The CUDA ecosystem has decades of hand-optimized kernel engineering, knowledge that is not available on newer platforms like Apple Silicon. With the growing adoption of MLX for local inference on MacBooks and Mac Studios (hundreds of millions of devices), closing this performance gap is crucial for developers to efficiently run language models and other AI workloads without relying on NVIDIA GPUs. This work demonstrates that it is possible to transfer that knowledge automatically, which could significantly accelerate the adoption of alternative hardware in AI.

Historically, porting kernels between ecosystems has been a huge challenge. For example, when AMD launched ROCm, the community had to manually rewrite many CUDA optimizations, a process that took years and still hasn't achieved full parity. Similarly, projects like OpenAI Triton have attempted to abstract kernel writing, but still require hardware-specific tuning. This work from Berkeley and IBM represents a qualitative leap by automating knowledge transfer, which could be a game-changer for the adoption of alternative hardware in AI.

How does the translation work?

The translation layer does not copy instruction by instruction, but maps CUDA optimization strategies to native MLX techniques. For example, CUDA shared memory optimizations are translated to the use of Apple Silicon's unified memory, and parallelization patterns are adapted to the architecture of Apple's GPU cores. The K-Search framework uses AI-guided evolutionary search to explore the kernel configuration space, starting from the translated CUDA kernels as seeds, and finds optimized variants for MLX.

Specifically, the process begins with an existing CUDA kernel, such as FlashAttention. The translation layer analyzes it and extracts key optimizations, such as using shared memory to reduce latency or tiling to improve data locality. It then generates an MLX kernel that implements those same strategies, but using native Metal primitives and unified memory. This translated kernel serves as a starting point for K-Search, which applies mutations and evolutionary selection to explore variants, evaluating their performance on real Apple Silicon hardware. The result is an optimized kernel that, in the case of FlashAttention, achieves 97% of the performance of the native MLX kernel manually developed by the community.

Consequences for the ecosystem

This research has profound implications:

• For AI developers: they will be able to run optimized models on Apple hardware without waiting for the community to develop native kernels. This lowers the barrier to experimenting with cutting-edge models on affordable and widely available hardware.
• For Apple: it strengthens MLX's position as a serious platform for AI inference, competing directly with CUDA. With the backing of an automated tool like K-Search, Apple could attract more developers to its ecosystem, especially those looking for alternatives to NVIDIA.
• For the industry: the approach is generalizable to other ecosystems (such as AMD ROCm or Google TPU), which could democratize access to high-performance kernels. Companies like AMD could benefit from a similar tool to automatically port CUDA kernels to ROCm, accelerating the adoption of their GPUs in AI.

Moreover, this advancement could reduce dependence on NVIDIA in the AI market, fostering greater competition and hardware diversity. For companies investing in AI infrastructure, this translates to more options and potentially lower costs.

Limitations and caveats

The authors note that the transfer is not perfect: they achieve 97% of native performance in FlashAttention, but there is a marginal loss. Additionally, the technique has been primarily evaluated on attention kernels and state space models; its generalization to other types of kernels (convolutions, normalizations) remains to be seen. The K-Search framework requires computational resources for evolutionary search, though it is a one-time cost per kernel. In their experiments, the authors used a cluster of NVIDIA GPUs to run the search, which could be a barrier for some developers, although the cost is amortized if the kernel is widely used.

Another limitation is that automatic translation may not capture all architecture-specific optimizations of Apple Silicon, such as the use of the Neural Engine or Apple's matrix units. The resulting kernels are competitive, but still do not outperform native ones in all cases. The MLX community will need to validate these results on a wider range of models and hardware.

What readers should know

This advancement is not an immediate replacement for manual optimization work, but a tool that drastically reduces the time and effort needed to port kernels between ecosystems. Developers working with MLX should watch for the integration of K-Search into tools like MLX Community or Hugging Face frameworks. For companies investing in Apple hardware for inference, this means they will soon be able to run models with performance close to NVIDIA's without rewriting code from scratch.

In a broader context, this work aligns with the trend toward automated compilation and optimization tools, such as XLA or TVM, which aim to abstract the complexities of underlying hardware. However, K-Search goes a step further by transferring expert knowledge between ecosystems, which could accelerate convergence toward a standard for portable kernels. As lead author Shiyi Cao notes: "Transferring a kernel from one vendor's hardware to another is harder than it looks, and often means rediscovering the same optimizations from scratch. This work shows that knowledge can be transferred automatically."

"Transferring a kernel from one vendor's hardware to another is harder than it looks, and often means rediscovering the same optimizations from scratch. This work shows that knowledge can be transferred automatically." — Shiyi Cao, lead author.