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

Researchers at Berkeley Sky Lab have extended K-Search, an AI-driven kernel optimization framework, with a structured translation layer that converts hard-won CUDA optimization knowledge into high-performance MLX kernels for Apple Silicon. The result: an attention kernel reaching 0.97x the speed of Apple's native implementation, and a Mamba SSM kernel delivering up to 20x faster prefill than the community mlx-lm version. The work addresses a growing bottleneck as AI hardware diversifies beyond NVIDIA.

Why the MLX Performance Gap Exists

MLX, Apple's machine-learning framework, has grown rapidly since late 2023. With Apple Silicon in hundreds of millions of MacBooks and Mac Studios, it enables local AI inference without cloud costs, and its unified memory architecture suits mid-sized models in the 7B–70B parameter range on M-series chips.

But many kernels the NVIDIA ecosystem takes for granted — paged attention, optimized SSM scan kernels, fused MoE routing — are missing or naive in MLX. Models run correctly, yet significant performance goes unused. Meanwhile, the CUDA ecosystem holds decades of hand-tuned expertise in attention, state space models, and other operations, representing thousands of engineering hours. Porting that knowledge manually to each new architecture has always meant rediscovering optimizations from scratch.

How K-Search Works

K-Search is an evolutionary kernel optimization framework originally developed by first author Shiyi Cao at UC Berkeley Sky Lab. Given a naive kernel and a hardware specification, it runs an iterative loop: a large language model reasons about which optimization to try next, generates a candidate kernel, then compiles and benchmarks it on real hardware. Measurements feed back into the search until performance converges.

The framework's persistent reasoning state, called a world model, is structured as a decision tree. Each root-to-leaf path composes a full optimization plan; every node carries an overall rating from 0 to 10, a confidence score from 0 to 1, and estimated impacts on memory bandwidth, register pressure, and hardware fit. When the best score stalls for several rounds, the search backs off and explores alternative branches. In these runs, a single model — Gemini 3.5 Pro Preview — handled both reasoning and code generation.

In its original evaluation on FlashInfer CUDA kernels (GQA decode, MLA decode, MLA prefill, and MoE), K-Search improved more consistently than OpenEvolve and ShinkaEvolve over the same 120-iteration budget.

Building the CUDA-to-MLX Bridge

The core challenge: simply handing an LLM a CUDA kernel and asking it to port produces syntactically valid but architecturally wrong code — wrong tile sizes, invalid primitives, mismatched memory assumptions. The team's translation layer provides three components:

• Concept mapping tables with hard constraints. For example, `__shared__` maps to Metal threadgroup memory but with a 32 KB limit versus NVIDIA's 48 KB, and H100's ~3.35 TB/s HBM3 corresponds to M3 Max's ~400 GB/s unified DRAM — a bandwidth gap that reshapes which optimizations matter.
• MLX-specific patterns for operations without CUDA equivalents, such as register-based row reductions using `simd_shuffle_xor` in an 8×8 MMA tile, and the "exp2 trick" — replacing exp(x) with exp2(x log₂ e) — to exploit Apple's fast exp2 hardware instruction.
• Reusable assertions that reframe expert kernel behaviors as properties the search must preserve rather than code to copy.

The backend also includes MLX-specific task adapters for compilation, generation prompts, and benchmarking via mlx.core utilities.

Attention and Mamba Results

Three attention kernel configurations isolate the translation layer's impact: a naive baseline, pure evolution without context, and full context. Pure evolution reached only 0.26x Apple's native attention kernel speed; with full context, it hit 0.97x. The fully informed kernel independently rediscovered key FlashAttention-2 strategies — threadgroup memory tiling, online softmax, K-transposition, double buffering, loop unrolling, and the exp2 trick.

On the Mamba SSM kernel (mamba-370m f16, M1 Max 64GB), the evolved implementation posted 152 tok/s decode versus mlx-lm's 116 and mamba.py's 40. Prefill results were more dramatic: 5,751 tok/s at L=512 versus 329 (mlx-lm) and 1,089 (mamba.py); 6,010 tok/s at L=1024 versus 327 and 1,127; and 6,612 tok/s at L=2048 — roughly 20x the community implementation.

What Comes Next

The team notes the method is not MLX-specific: any ecosystem where CUDA expertise is transferable — custom AI accelerators included — could use the same approach. As non-NVIDIA hardware keeps multiplying, automated knowledge transfer between kernel ecosystems may become a standard tool for closing performance gaps that would otherwise take years of manual engineering.