From CUDA to MLX: How K-Search Brings Decades of Kernel Expertise to Apple Silicon

Every fast AI model eventually comes down to some engineer's hand-tuned GPU kernel, the tiny bit of low-level code that decides whether your hardware hums or idles. NVIDIA's CUDA ecosystem has thousands of these, refined over more than a decade. Apple Silicon, despite running on hundreds of millions of Macs and being a genuinely appealing platform for local inference thanks to its unified memory, has almost none of that inherited wisdom. MLX, Apple's ML framework, runs models correctly but often leaves serious performance sitting on the table because nobody has spent the CUDA-equivalent of a decade tuning its kernels.

A team building on K-Search, an evolutionary kernel optimizer out of Berkeley's Sky Lab, decided to test whether that CUDA expertise could just be translated rather than rediscovered. K-Search already works by having an LLM propose optimizations, generate code, benchmark it on real hardware, and iterate — pruning dead ends and refining promising branches in something like a decision tree. The team's contribution here was a structured CUDA-to-MLX translation layer: concept-mapping tables that tell the model, for instance, that Metal's threadgroup memory caps out at 32 KB versus CUDA's 48 KB, or that an H100's 3.35 TB/s of bandwidth becomes roughly 400 GB/s on an M3 Max, a gap large enough to flip which optimizations are even worth attempting.

The payoff showed up clearly on an attention kernel. Handed no context, the evolved kernel limped along at 0.26x the speed of Apple's native attention implementation. Give it the translation layer, though, and it climbed to 0.97x — essentially matching hand-tuned performance — by independently rediscovering FlashAttention-2's playbook: threadgroup tiling, online softmax, and a neat trick where exponentials get rewritten in base 2 so the kernel can hit Apple's fast exp2 hardware instruction directly.

The more dramatic number came from a Mamba state-space model kernel. Tested on an M1 Max, the evolved kernel hit 6,600-plus tokens per second on long-sequence prefill, versus around 330 tokens per second for the community mlx-lm implementation — a roughly 20x gap. The reason is almost embarrassingly simple: mlx-lm processes the SSM's recurrence one token at a time, even though the underlying math is associative and can be restructured into a parallel scan that finishes in log(N) steps instead of N. The evolved kernel found that scan; the community one never bothered, leaving most of the GPU idle during prefill. Decode speeds, where there's only one token to process anyway, stayed roughly comparable across implementations, which is exactly what you'd expect if parallel scanning is the whole story.

The researchers are careful to note this is two kernels, not a proof that every optimization transfers cleanly across silicon vendors. But the framing is the interesting part: the limiting factor wasn't whether an LLM could write competent Metal code. It was whether it had been given the right hardware-specific context and constraints to reason with. That's a fairly different problem than the one most kernel-generation efforts have been solving.