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

What it claims

The page attributes most of the attention result to the CUDA-to-MLX translation layer, not to the search itself. On Apple's MLX attention kernel, the evolved kernel goes from 0.26x to 0.97x of native speed when the translation context is supplied. The roughly 20x Mamba gain over mlx-lm has a specific source: mlx-lm runs the SSM recurrence token by token, and the evolved Metal kernel uses a parallel prefix scan instead. The authors conclude that the limiting factor was "the quality of the context and constraints" given to the LLM, not the LLM's ability to write Metal.

How it works

• Reasoning step. Before proposing anything, the reasoner works through a fixed analysis: 1. Classify the kernel (reduction, scan, attention/softmax, …). 2. Rewrite the reference computation in canonical form. 3. Map the data layout and access patterns. 4. Hypothesize the bottleneck (bandwidth, latency, compute or synchronization) for each runtime regime.
• Candidate format. Each candidate is a single change that can be implemented in one iteration.
• World model. The persistent state is a prefix tree. Each root-to-leaf path is a full optimization plan, and siblings are competing alternatives. Every node carries:
• an overall_rating_0_to_10
• a confidence_0_to_1
• impact scores for memory bandwidth, register pressure and compute/hardware fit

Refining an idea adds a child node rather than overwriting the parent. If the best score stalls for a stagnation window, the search backs off to another branch.

• Translation layer. It has three parts:
• Concept tables with hard constraints. __shared__ maps to threadgroup memory, capped at 32 KB versus NVIDIA's 48 KB. warp_reduce maps to MMA. __syncthreads() maps to threadgroup_barrier(mem_flags::mem_tg). The H100's ~3.35 TB/s is set against the M3 Max's ~400 GB/s, a gap that changes which optimizations are worth pursuing.
• Metal-specific patterns. Examples are simd_shuffle_xor row reductions in an 8×8 MMA tile, and the exp2 trick (fast::exp2).
• Assertions. Expert kernel behaviours are written as properties the search must preserve, not as code to copy.

flowchart TD
  T[CUDA expert kernels] --> X[Translation layer]
  X --> R[Reasoner: classify, bottleneck]
  R --> W[World-model tree]
  W --> G[Generate one-change candidate]
  G --> B[Compile + benchmark on Mac]
  B --> S{Best score improved?}
  S -- yes --> W
  S -- stalled --> A[Back off to sibling branch] --> W

Evidence

Attention. Three configurations were compared: a naive baseline, pure evolution with no added context, and the full translation layer. The full-context result is 0.97x of Apple's native MLX attention, and the page's comparison point for the translation layer is 0.26x. With full context, the search independently found the FlashAttention-2 techniques: threadgroup tiling, online softmax, K-transposition and exp2.

Mamba SSM. The setting is mamba-370m in f16 on an M1 Max with 64 GB. Prefill is about 20x faster than mlx-lm. Decode is described as "roughly flat", because a single new token gives the scan nothing to parallelize. mamba.py is slow on both because it falls back to CPU or MPS.

Original K-Search paper. On FlashInfer CUDA kernels (GQA decode, MLA decode, MLA prefill, MoE) with a 120-iteration budget, K-Search improved more consistently than OpenEvolve and ShinkaEvolve.

How to use it

git clone https://github.com/caoshiyi/K-Search.git
cd K-Search
uv pip install openai wandb
uv pip install git+https://github.com/caoshiyi/flashinfer-bench-ksearch.git

Set KSEARCH_ROOT and API_KEY at the top of the script under scripts/, then run one of:

bash scripts/mac_flash_attention_wm.sh
bash scripts/mamba_selective_scan_fwd_wm.sh

The search runs used Gemini 3.5 Pro Preview for both reasoning and code generation.

Limits

• Narrow evaluation. Only two kernels were studied, and the authors say they "do not yet know how far this generalizes."
• Hardware unclear for attention. The Mac model used for the attention benchmark is not stated (M3 Max appears only in the mapping table). The page's text gives no attention timings and no decode figures beyond "roughly flat".
• Ablation label unclear. The page does not say which no-context configuration produced the 0.26x figure.
• Not yet covered. Paged attention and fused MoE routing are future work.
• Setup instructions inconsistent. The setup step says to set "three variables" but lists only two.
• Licence. No licence is stated.