Glonce

K-Search translates CUDA kernel expertise for Apple's MLX

researchThe Glonce editorsBerkeley AI Research🔥68Glonce rating68Impact68Novelty62Relevance85Surprise48Players52Credibility85
K-Search translates CUDA kernel expertise for Apple's MLX

GPU kernels, the low-level programs that run inside a GPU, are a crucial and hard-to-write component of AI systems. The CUDA ecosystem has accumulated decades of hand-tuned kernel expertise for attention, state space models and other operations, representing thousands of engineering hours. Newer hardware ecosystems, including Apple Silicon, are growing fast but lack that depth. Apple's MLX framework has seen strong adoption since late 2023: with Apple Silicon in hundreds of millions of MacBooks and Mac Studios, MLX enables local AI inference without cloud costs, and its unified memory architecture suits mid-sized models in the 7B to 70B parameter range on M-series chips. Despite that momentum, MLX still lacks tuned versions of several performance-critical kernels the NVIDIA ecosystem takes for granted, including paged attention, optimized SSM scan kernels and fused MoE routing; it runs models correctly but leaves performance on the table.

The post builds on K-Search, an evolutionary kernel optimization framework introduced by Cao et al. at Berkeley Sky Lab. Given a naive kernel and a hardware specification, K-Search runs an iterative loop: an LLM reasons about which optimization to try next, a code-writing model generates candidate kernels, and those candidates are compiled and benchmarked on real hardware, with measurements feeding back into a persistent 'world model', a scored decision tree of candidate optimizations that keeps growing across rounds. The original K-Search paper evaluated this approach on CUDA kernels from FlashInfer (GQA decode, MLA decode, MLA prefill, MoE) and reported that it improved more consistently than OpenEvolve and ShinkaEvolve over the same 120-iteration budget.

This work adds a native MLX backend to K-Search, covering kernel compilation and execution via MLX's Metal/C++ APIs, updated kernel-generator prompts for Metal/MLX, and benchmarking integration using mlx.core. The more interesting contribution is a structured CUDA-to-MLX translation layer that lets K-Search treat existing CUDA kernels as a knowledge base rather than rebuilding from scratch, because simply handing an LLM a CUDA kernel and asking it to port it produces code that compiles but is architecturally wrong. The layer has three parts: concept mapping tables that pair CUDA primitives with MLX/Metal equivalents and hard constraints (for example, CUDA's shared maps to Metal threadgroup memory but with a 32 KB limit versus NVIDIA's 48 KB; warp_reduce maps to MMA; __syncthreads() becomes threadgroup_barrier(mem_flags::mem_tg); and the H100's roughly 3.35 TB/s HBM3 bandwidth maps to the M3 Max's roughly 400 GB/s unified DRAM, a gap that reshapes which optimizations are worth pursuing); MLX-specific hints and patterns, such as register-based row reductions using simd_shuffle_xor in an 8x8 MMA tile layout, and the 'exp2 trick' that replaces exp(x) with 2 to the power of x times log2(e) to use Apple's fast exp2 hardware instruction instead of paying for a base conversion; and reusable assertions that reframe expert kernel behaviors as properties the search must preserve rather than code to copy. In the authors' runs, a single model, Gemini 3.5 Pro Preview, plays both the reasoning role and the code-writing role.

On the attention kernel, the authors compared a naive baseline, pure evolution with no added context, and the full translation layer. The evolved kernel jumped from 0.26x to 0.97x the speed of Apple's own state-of-the-art MLX Attention kernel once the translation layer supplied context. With full context, the search independently rediscovered the core FlashAttention-2 optimizations: threadgroup memory tiling, online softmax, K-transposition for memory access, and the exp2 trick.

To test whether the approach generalizes beyond attention, the authors applied it to the Mamba state-space model kernel, whose bottleneck is a recurrent state update rather than a softmax. They compared their evolved kernel against the community mlx-lm implementation and the PyTorch reference mamba.py, running mamba-370m in f16 on an M1 Max with 64 GB of memory. The evolved kernel reached up to a 20x faster prefill than mlx-lm, and on decode it hit 152 tokens per second against 116 tokens per second for mlx-lm.

Key facts

  • K-Search, an evolutionary GPU kernel search framework from Cao et al. at Berkeley Sky Lab, gained a new backend for Apple's MLX framework running on Apple Silicon.
  • A structured CUDA-to-MLX translation layer (concept mapping tables, MLX-specific hints, reusable assertions) lets K-Search reuse existing CUDA kernel expertise instead of rebuilding from scratch.
  • With the full translation layer, the evolved MLX attention kernel reached 0.97x the speed of Apple's native MLX Attention kernel, up from 0.26x with no added context.
  • The evolved Mamba SSM kernel reached up to a 20x faster prefill than the community mlx-lm implementation, and 152 tokens per second on decode versus 116 for mlx-lm, on mamba-370m f16 on an M1 Max with 64 GB.
  • A single model, Gemini 3.5 Pro Preview, served as both the reasoning engine and the code-generation model throughout the search.

Why it matters

The post shows that decades of CUDA kernel optimization knowledge can be transferred automatically to a different hardware ecosystem instead of being rediscovered by hand. Apple's MLX has grown fast since late 2023 as a way to run local AI inference on Apple Silicon, but it still lacks tuned versions of performance-critical kernels like paged attention and SSM scans that the CUDA ecosystem already has. A translation layer that carries that expertise over closes part of that gap without requiring engineers to relearn every optimization from first principles on the new hardware.

Who it affects

Most directly, developers running local AI inference on Apple Silicon Macs through MLX, particularly for mid-sized models in the 7B to 70B parameter range that benefit from MLX's unified memory architecture. More broadly, it is relevant to anyone maintaining kernel libraries for non-NVIDIA AI hardware, since the authors state the translation approach is not specific to MLX and applies to any ecosystem where CUDA expertise could transfer.

How to use it

This is a research result rather than a shipped product: the post describes an MLX backend and a CUDA-to-MLX translation layer added to the existing K-Search framework, evaluated on an attention kernel and a Mamba SSM kernel. No pricing, licensing or release details for the code are given in the text.

How solid is it

The authors report concrete before/after numbers rather than only qualitative claims: the attention kernel moved from 0.26x to 0.97x the speed of Apple's native MLX Attention kernel once the translation layer supplied context, and independently rediscovered known FlashAttention-2 techniques (threadgroup memory tiling, online softmax, K-transposition, the exp2 trick). The Mamba SSM kernel reached up to 20x faster prefill and 152 tokens per second on decode versus 116 for the community mlx-lm implementation on mamba-370m f16 on an M1 Max. The underlying K-Search framework itself was previously benchmarked against OpenEvolve and ShinkaEvolve on CUDA FlashInfer kernels under a matched 120-iteration budget.

Risks and caveats

The reported numbers come from the authors' own benchmarks on specific kernel types (attention, Mamba SSM) and specific hardware (M1 Max), so it is not established how the approach performs on other kernel classes or other Apple Silicon chips. The source text is cut off partway through the decode benchmark table, so the mamba.py reference figure for that comparison is not available here. Results also depend on a single model, Gemini 3.5 Pro Preview, handling both the reasoning and code-generation roles, which ties outcomes to that model's specific capabilities.