K-Search, an evolutionary AI-driven kernel optimization framework from UC Berkeley Sky Lab, has been extended with an MLX backend to automatically transfer CUDA kernel expertise to Apple Silicon. The key innovation is a structured CUDA-to-MLX translation layer that maps CUDA primitives to Metal/MLX equivalents with hardware-specific constraints, enabling an LLM-guided evolutionary search to discover high-quality GPU kernels without GPU experts starting from scratch. Results show the evolved attention kernel reaches 0.97x the speed of Apple's native MLX attention kernel, and the Mamba SSM kernel achieves up to 20x prefill speedup over the community mlx-lm implementation by implementing a parallel prefix scan for the state-space model recurrence. The work is from IBM Research building on the UC Berkeley Sky Lab K-Search framework, and the MLX backend is open source and reproducible.
