From CUDA to MLX: K-Search Expertise Applied to Apple Silicon

A New Era in Computing

The computing landscape is rapidly evolving. Hardware advancements go beyond faster GPUs, introducing a variety of chips from different vendors, each with unique architectures tailored for specific AI workloads. Simultaneously, software development is accelerating, with AI coding tools capable of generating in minutes what previously took months.

As AI becomes central to computing, GPU kernels are essential. These low-level programs run inside GPUs and require significant expertise to optimize effectively. Transferring a kernel from one vendor’s hardware to another presents an even greater challenge, often requiring a rediscovery of optimizations specific to each architecture.

The Expertise of CUDA Kernels

The CUDA ecosystem has accumulated decades of research and optimizations, leading to highly efficient implementations of critical functions like attention and state space models. These achievements represent thousands of engineering hours. However, newer hardware ecosystems, such as Apple Silicon and custom AI accelerators, are growing quickly but lack this depth of expertise.

It becomes crucial to ask whether this expertise can be transferred automatically.

K-Search: A Leap in Kernel Optimization

K-Search, an evolutionary kernel optimization framework developed at UC Berkeley, addresses this challenge. This method utilizes artificial intelligence to optimize GPU kernels. By extending K-Search with a backend for MLX, Apple’s machine learning framework, we’ve created a CUDA-to-MLX translation layer. This approach allows K-Search to leverage existing CUDA kernels to generate new ones for Apple Silicon, thereby reducing the need to rebuild kernels from scratch.

The results are impressive: our method achieves near-expert performance on Apple Silicon, offering a speedup of 0.97x compared to the native MLX attention kernel and up to 20x in prefill speed over the community mlx-lm implementation on the Mamba SSM kernel.

Why Choose MLX?

Since late 2023, Apple’s MLX framework has seen remarkable adoption. With millions of MacBooks and Mac Studios equipped with Apple Silicon, MLX facilitates local AI inference without cloud costs. Its unified memory architecture makes it particularly suitable for mid-sized models (7B to 70B parameters on M series chips).

However, a significant gap remains: many performance-critical kernels commonly found in the NVIDIA ecosystem are either absent or under-optimized on MLX without hardware-specific tuning. While MLX executes models correctly, it often leaves considerable performance on the table.

How K-Search Works

K-Search distinguishes itself by iterating on a naive kernel based on hardware specifications. An AI model reasons about optimizations to apply, generates candidate kernels, which are then compiled and tested on real hardware. The resulting measurements feed back into the search process, allowing continuous improvement until performance converges.

This process is guided by a domain-specific document that encodes hardware rules and mathematical constraints, preventing the generation of invalid code and ensuring candidates compile and run efficiently. In practice, a single model (like Gemini 3.5 Pro Preview) performs both reasoning and code generation by analyzing the kernel before proposing optimizations.

Conclusion

In summary, the combination of accumulated expertise in the CUDA ecosystem and the innovation brought by K-Search opens new avenues for performance optimization on Apple Silicon. The applications of this method are not limited to MLX but extend to any ecosystem where CUDA expertise can be beneficially transferred.

If you want to learn more about optimizing the performance of your AI infrastructure or implementing similar strategies, feel free to Contact me.