从 CUDA 到 MLX：K-Search 如何将数十年 Kernel 经验带到 Apple Silicon

图 1：CUDA 到 MLX 的优化翻译图谱。 CUDA 优化知识可以被转化为适配特定架构的 MLX 策略，而非逐条指令照搬。

我们正处在计算的新纪元。硬件正在快速变化——不仅是更快的 GPU，还有来自不同厂商、各具架构且通常针对特定 AI 工作负载定制的不断扩大的芯片品类。软件变化同样迅速，AI 编程工具如今几分钟就能生成数年前需要数月才能完成的代码。

随着计算越来越以 AI 为中心，GPU Kernel 已成为其成功的关键组成部分。这些是在 GPU 内部运行的底层程序，编写高效 Kernel 绝非易事——需要多年的专业经验才能做好。将 Kernel 从一家厂商的硬件迁移至另一家的硬件则更困难，往往意味着要从头重新发现相同的优化方法。例如，CUDA 生态已经积累了数十年艰难获得的 Kernel 专业知识：针对 Attention、状态空间模型及其他关键操作的手工调优实现，代表着数千小时的工程投入。较新的硬件生态（Apple Silicon、定制 AI 加速器等）虽然发展迅速，却缺乏这样的深度积累。

在这项工作中，我们探讨这种专业知识能否被自动迁移。我们基于 Berkeley Sky Lab 的 Cao 等人提出的 K-Search——一个利用 AI 优化 GPU Kernel 的进化式 Kernel 搜索框架——进行开发，并为其扩展了 MLX 后端；MLX 是 Apple 面向其 Apple Silicon 芯片推出的机器学习框架。我们开发了一套新颖的结构化 CUDA-to-MLX 翻译层，使 K-Search 能够将现有 CUDA Kernel 作为知识库，把它们适配为 Apple Silicon 上的高质量 GPU Kernel，而不是从零开始重建。

我们表明，该方法在 Apple Silicon 上实现了接近专家水平的性能：相对于原生 MLX Attention Kernel，达到 0.97× 的速度；在 Mamba SSM Kernel 上，相对于社区 mlx-lm 实现，prefill 速度最高可提升 20×。下文将报告具体数字，以及其中有多少收益来自翻译层。尽管我们聚焦于 Apple Silicon 的 MLX Kernel，但该方法并不局限于 MLX，可应用于任何能够迁移 CUDA 专业知识的生态系统。

为什么选择 MLX？

自 2023 年末以来，Apple 的 MLX 框架获得了显著采用。随着数亿台 MacBook 和 Mac Studio 搭载 Apple Silicon，MLX 让本地 AI 推理无需承担云端成本。统一内存架构使其特别适合处理中等规模模型（M 系列芯片上可运行 7B–70B 参数模型）。

然而，在这一发展势头之下仍存在明显缺口：许多 NVIDIA 生态中被视为理所当然、且对性能至关重要的 Kernel——如分页 Attention、优化的 SSM scan Kernel、融合式 MoE 路由——在 MLX 中要么缺失，要么没有针对硬件进行专门调优而显得较为朴素。MLX 能够正确运行模型，但往往留下了大量未被释放的性能空间。

这一缺口正是本文其余部分的出发点。

什么是 K-Search？

K-Search 是一个进化式 Kernel 优化框架，最初由我们的第一作者、UC Berkeley Sky Lab 的 Shiyi Cao 开发。给定一个朴素 Kernel 和硬件规格，它会运行一个迭代优化循环：LLM 推理下一步应尝试哪些优化；代码生成模型编写候选 Kernel；随后在真实硬件上编译并测试这些候选方案。

测量结果会反馈到搜索过程，后者持续细化方案、探索有前景的方向并放弃无效路径，直至性能收敛。

算法 1：通过协同进化世界模型实现 K-Search。 搜索过程在选择最有前景的动作、实例化并评估代码直至改进停滞，以及通过插入、更新和剪枝操作进化世界模型之间交替进行。改编自 Cao 等人（2026）。

搜索由 Spec 进行约束：这是一份领域特定文档，编码了硬件规则、优化模式和数学约束。它可避免生成的代码臆造无效原语，并确保候选方案能够真正编译并高效运行。

在我们的实验中，单一模型（Gemini 3.5 Pro Preview）同时承担两种角色：维护推理状态和编写 Kernel。推理部分被提示为一名“GPU Kernel 性能工程师”，并被要求在提出任何方案前完成固定分析：对 Kernel 进行分类（归约、扫描、attention/softmax 等）、将参考计算重写为规范形式、梳理数据布局和访问模式，以及针对每种运行时情况推测可能的瓶颈（带宽、延迟、计算或同步）。只有完成这些步骤后，它才会输出候选优化，每项均为可在一次迭代中实现的单一改动。

我们将这种持久化推理状态称为 world model（世界模型）。它不是一份扁平的待尝试事项列表，而是一棵决策（前缀）树：每条从根节点到叶节点的路径都构成一套完整的优化方案，而同级分支则是相互竞争的替代方案。每个节点都会被评分——包括范围为 [0, 10] 的 overall_rating、范围为 [0, 1] 的 confidence，以及针对内存带宽、寄存器压力和计算/硬件适配度的节点级 impacts——因此搜索可以为部分方案排序，并扩展其中最有前景的方案。这棵树会跨轮次持久保存并持续生长：细化一个想法会新增子节点，而不是覆盖其父节点；如果最佳分数在数轮内未能提升（即停滞窗口），搜索将回退并探索替代分支。Attention Kernel 的一次运行中，一个处于中间状态的节点如下所示：

{
  "action": "Replace the threadgroup-memory softmax reduction
             with a register-only reduction: each SIMD group
             owns 8 query rows and reduces across lanes with
             simd_shuffle_xor, removing a threadgroup_barrier.",
  "difficulty_1_to_5": 4,
  "impacts": {
    "memory_bandwidth":  8,
    "register_pressure": 4,   // risk: spill if Br > 8
    "compute_hw_fit":    9    // SIMD width 32; keep tile 8x8
  },
  "overall_rating_0_to_10": 8,
  "confidence_0_to_1": 0.7
}

代码清单 1：K-Search 世界模型节点示例。 每项候选优化均记录一项具体动作、预估硬件影响、整体优先级评分以及模型的置信度。

图 2：K-Search 概览。 该框架作用于一个以搜索树组织的搜索状态 $S_t$。该树由 Closed 节点（蓝色，带有附加程序、如 $x_{12}$ 的已访问状态）和 Open 节点构成的 Frontier（橙色，待处理假设、如 $u_{13}$）组成。工作流循环经过三个阶段：(1) 动作选择：基于世界模型估计的优先级分数 $V$，从 Frontier 中检索最有前景的动作节点；(2) 局部细化：随机策略 $\pi_{\mathrm{code}}$ 对具体实现进行采样，直至停滞；(3) 世界模型更新：LLM 根据轨迹推理，通过 Insert（插入新动作）、Update（调整 $V$，例如 $u_{11}$ 从 0.9 降至 0.6）和 Prune（移除较不具前景的节点，如 $u_{10}$）更新搜索树。

原始 K-Search 论文在 FlashInfer 的 CUDA Kernel 上评估了这一搜索策略。在相同的 120 次迭代预算下，针对 GQA decode、MLA decode、MLA prefill 和 MoE，K-Search 比 OpenEvolve 和 ShinkaEvolve 实现了更稳定的改进。这些结果确立了我们在此基础上构建的搜索框架；本文其余部分将探讨其优化知识能否迁移到 CUDA 之外。

图 3：原始 K-Search 论文的主要结果。 在三次运行中，针对四个 FlashInfer CUDA Kernel，K-Search 在当前最优搜索分数、各工作负载 Kernel 性能及加速比分布方面均优于 OpenEvolve 和 ShinkaEvolve。完全复现自 Cao 等人（2026）。

构建 MLX 后端

为了将 K-Search 引入 Apple Silicon，我们首先构建了原生 MLX 后端。我们为 K-Search 实现了一套完整的 MLX 专用任务适配器，包括：

• 位于 k_search/tasks/ 的 MLX 任务后端，通过 MLX 的 Metal/C++ API 处理 Apple Silicon 上的 Kernel 编译与执行。
• 用于编写和修改 Metal/MLX Kernel 的更新版 Kernel 生成器提示词。
• 使用 mlx.core 测量工具的 MLX 专用基准测试集成。

将 CUDA 专业知识翻译至 MLX

然而，更有意思的挑战并不只是让 K-Search 在 MLX 上运行。关键洞见在于，专家级 CUDA Kernel 编码了数十年的优化知识；只要能够弥合概念差距，这些知识就可以迁移至 Apple GPU。仅仅将一个 CUDA Kernel 交给 LLM 并要求其移植还远远不够：如果缺少深厚的硬件上下文，模型会生成语法上有效、但在架构层面错误的代码（例如错误的 tile 大小、无效原语和不匹配的内存假设）。

我们的翻译层包括：

• 概念映射表： 一份包含 CUDA 原语及其 MLX/Metal 等价物的结构化术语表，并附带硬性约束。例如：
  • __shared__ 映射至 Metal threadgroup 内存，但其硬性上限为 32 KB（而 NVIDIA 为 48 KB）
  • warp_reduce 映射至 MMA（优先选择）
  • __syncthreads() 变为 threadgroup_barrier(mem_flags::mem_tg)
  • H100 约 3.35 TB/s 的 HBM3 映射至 M3 Max 约 400 GB/s 的统一 DRAM；这种带宽差异会重塑哪些优化值得优先追求。
• MLX 专用提示与模式： 面向没有直接 CUDA 等价物的操作提供具体代码级模式，例如：在 8×8 MMA tile 布局中，使用 simd_shuffle_xor 实现基于寄存器的行归约；或采用“exp2 技巧”（以 $exp_2(x \log_2 e)$ 替换 $exp(x)$），利用 Apple 快速的 $exp_2$ 硬件指令来加速 softmax。
• 可复用断言： 将专家 Kernel 的行为重新表述为进化搜索必须保持的属性，而不是需要直接复制的代码。

匹配专家 Kernel 性能：Attention Kernel

我们评估了适用于 Apple Silicon 的 MLX Attention Kernel 的三种配置：(1) 朴素基线；(2) 未提供额外上下文的纯进化；以及 (3) 完整上下文翻译层。第三种配置为优化器提供了从高性能 Kernel（例如 FlashAttention-2）中提取的、特定于架构的实现知识，使进化搜索能够针对实现策略进行推理，而不必从朴素 Kernel 起步。这三种配置共同帮助我们分离出翻译层的确切影响。

图 4：通过叠加优化实现 Attention Kernel 性能扩展。 “完整上下文”配置成功发现并实现了双缓冲和循环展开等高级策略，达到了接近专家水平的性能。

从 Apple 最先进 Attention Kernel 速度的 0.26× 跃升至 0.97×，说明了翻译层的重要性。借助完整上下文，演化出的 Kernel 独立发现了 FlashAttention 2 的关键优化：threadgroup 内存分块、在线 softmax、用于内存访问的 K 转置，以及 exp2 技巧。其中最后一项将每个 softmax 指数运算替换为以 2 为底的指数运算：

\[e^x = 2^{x \log_2 e},\]

这在数学上是精确的，并使 Kernel 能够直接使用 Apple 快速的 fast::exp2() 硬件指令，而无需在运行时承担底数转换开销。

快 20× 的 prefill：Mamba SSM Kernel

为评估 K-Search 能否泛化至 Attention Kernel 之外，我们将其应用于 Mamba 使用的状态空间模型（SSM）Kernel。与 Attention 不同，其计算瓶颈是循环状态更新而非 softmax，因此提供了截然不同的优化挑战。我们在一台 M1 Max 上，将演化得到的实现与社区 MLX 实现（mlx-lm）及 PyTorch 参考实现（mamba.py）进行对比。

在 mamba-370m f16、M1 Max 64GB 上评估：

指标	mlx-mamba（我们的实现）	mlx-lm（社区实现）	mamba.py
Decode	152 tok/s	116 tok/s	40 tok/s
Prefill L=512	5,751 tok/s	329 tok/s	1,089 tok/s
Prefill L=1024	6,010 tok/s	327 tok/s	1,127 tok/s
Prefill L=2048	6,612 tok/s	326 tok/s	1,092 tok/s
Prefill L=4096	6,743 tok/s	339 tok/s	1,042 tok/s

表 1： mamba-370m（f16、M1 Max 64GB）上的 prefill 和 decode 吞吐量。mlx-mamba（我们的实现）实现了比社区 mlx-lm 基线高约 20× 的 prefill 吞吐量，同时 decode 性能保持相近。

相对于 mlx-lm 约 20× 的 prefill 加速，归根结底来自一个差异：mlx-lm 没有为 SSM 实现并行扫描。状态递推关系

\[h_t = \bar{a}t h{t-1} + \bar{b}_t\]

看起来本质上是串行的，但每一步都可以写成一对 $(\bar{a}_t, \bar{b}_t)$，并在如下结合运算下进行组合：

\[(a_2, b_2) \circ (a_1, b_1) = \left(a_2 a_1,\ a_2 b_1 + b_2\right),\]

该运算可精确复现上述递推关系。由于该算子具有结合性，整个序列可通过并行（前缀）扫描在 $O(\log N)$ 的深度下完成

Figure 1: CUDA-to-MLX optimization translation map. CUDA optimization knowledge can be translated into architecture-native MLX strategies rather than copied instruction-for-instruction.

We face a new epoch in computing. Hardware is changing rapidly — not just faster GPUs, but a growing range of chips from different vendors, each with its own architecture and often tailored to specific AI workloads. Software is changing just as fast, and AI coding tools now generate in minutes what took months of effort a few years ago.

With so much of computing now centered on AI, GPU kernels are a crucial component of its success. These are the low-level programs that run inside the GPU, and writing efficient ones is far from obvious — it takes years of expertise to get right. Transferring a kernel from one vendor’s hardware to another is harder still, and often means rediscovering the same optimizations from scratch. The CUDA ecosystem, for example, has accumulated decades of hard-won kernel expertise: hand-tuned implementations of attention, state space models, and other critical operations representing thousands of engineering hours. Newer hardware ecosystems (Apple Silicon, custom AI accelerators, and others) are growing fast but lack this depth.

In this work we ask whether that expertise can be transferred automatically. We built on K-Search, an evolutionary kernel search framework introduced by Cao et al. at Berkeley Sky Lab that uses AI to optimize GPU kernels, and extended it with a backend for MLX — Apple’s machine-learning framework for its own Apple Silicon chips. We developed a novel structured CUDA-to-MLX translation layer that lets K-Search take existing CUDA kernels as a knowledge base and adapt them into high-quality GPU kernels for Apple Silicon, rather than rebuilding from scratch.

We show that our approach reaches near-expert level performance on Apple Silicon with 0.97x speedup compared to the native MLX Attention kernel, and up to a 20x prefill speedup over the community mlx-lm implementation on the Mamba SSM kernel; we report the numbers, and how much of the gain comes from the translation layer, in the sections below. Although we focus on MLX kernels for Apple Silicon, the method is not specific to MLX and applies to any ecosystem where CUDA expertise is transferable.

Why MLX?

Apple’s MLX framework has seen remarkable adoption since late 2023. With Apple Silicon in hundreds of millions of MacBooks and Mac Studios, MLX enables local AI inference without cloud costs. The unified memory architecture makes it especially attractive for mid-sized models (7B–70B parameters on M series chips).

Yet beneath this momentum lies a significant gap: many performance-critical kernels that the NVIDIA ecosystem takes for granted: paged attention, optimized SSM scan kernels, fused MoE routing are either absent or naive without hardware-specific tuning. MLX runs models correctly but often leaves significant performance on the table.

This gap is what motivates the rest of this post.

What is K-Search?

K-Search is an evolutionary kernel optimization framework originally developed by our first author Shiyi Cao at UC Berkeley Sky Lab. Given a naive kernel and a hardware specification, it runs an iterative optimization loop: an LLM reasons about which optimizations to try next, a code-writing model generates candidate kernels, and those candidates are compiled and benchmarked on real hardware.

Measurements feed back into the search, which keeps refining, pursuing promising directions and dropping dead ends until performance converges.

Algorithm 1: K-Search via co-evolving world models. The search alternates between selecting the most promising action, instantiating and evaluating code until improvement stagnates, and evolving the world model through insert, update, and prune operations. Adapted from Cao et al. (2026).

Search is grounded by a Spec: a domain-specific document encoding hardware rules, optimization patterns, and mathematical constraints which keeps generated code from hallucinating invalid primitives and ensures candidates will actually compile and run efficiently.

In our runs, a single model (Gemini 3.5 Pro Preview) plays both roles: it maintains the reasoning state and writes the kernels. The reasoning half is prompted as a “GPU kernel performance engineer” and asked to work through a fixed analysis before proposing anything: classify the kernel (reduction, scan, attention/softmax, …), rewrite the reference computation in canonical form, map out data layout and access patterns, and hypothesize the likely bottleneck (bandwidth, latency, compute, or synchronization) in each runtime regime. Only then does it emit candidate optimizations, each as a single change implementable in one iteration.

We call the persistent reasoning state a world model. Rather than a flat list of things to try, it is a decision (prefix) tree: each root→leaf path composes a full optimization plan, and sibling branches are competing alternatives. Every node is scored — an overall_rating in [0, 10], a confidence in [0, 1], and per-node impacts on memory bandwidth, register pressure, and compute/hardware fit — so the search can rank partial plans and expand the most promising ones. The tree persists and grows across rounds: refining an idea adds a child node rather than overwriting its parent, and if the best score fails to improve for a few rounds (a stagnation window) the search backs off to explore an alternative branch. A single node, as it appears mid-run on the attention kernel, looks like this:

{
  "action": "Replace the threadgroup-memory softmax reduction
             with a register-only reduction: each SIMD group
             owns 8 query rows and reduces across lanes with
             simd_shuffle_xor, removing a threadgroup_barrier.",
  "difficulty_1_to_5": 4,
  "impacts": {
    "memory_bandwidth":  8,
    "register_pressure": 4,   // risk: spill if Br > 8
    "compute_hw_fit":    9    // SIMD width 32; keep tile 8x8
  },
  "overall_rating_0_to_10": 8,
  "confidence_0_to_1": 0.7
}

Listing 1: Example K-Search world-model node. Each candidate optimization records a concrete action, estimated hardware impacts, an overall priority rating, and the model's confidence.

Figure 2: Overview of K-Search. The framework operates on a Search State $S_t$ structured as a search tree. The tree consists of Closed nodes (blue, visited states with attached program like $x_{12}$) and a Frontier of Open nodes (orange, pending hypotheses like $u_{13}$). The workflow iterates through three phases: (1) Action Selection , where the most promising action node is retrieved from the frontier based on world model estimated priority score $V$; (2) Local Refinement , where a stochastic policy $\pi_{\mathrm{code}}$ samples concrete implementations until stagnation; and (3) World Model Update , where the LLM reasons over the trajectory to update the search tree via Insert (adding new actions), Update (adjusting $V$, e.g., $u_{11}$ dropping from 0.9 to 0.6), and Prune (removing less promising nodes like $u_{10}$).

The original K-Search paper evaluated this search strategy on CUDA kernels from FlashInfer. Across GQA decode, MLA decode, MLA prefill, and MoE, K-Search improved more consistently than OpenEvolve and ShinkaEvolve over the same 120-iteration budget. These results establish the search framework we build on here; the remainder of this post asks whether its optimization knowledge can transfer beyond CUDA.

Figure 3: Main results from the original K-Search paper. Across three runs, K-Search achieves stronger best-so-far search scores, per-workload kernel performance, and speedup distributions than OpenEvolve and ShinkaEvolve on four FlashInfer CUDA kernels. Reproduced exactly from Cao et al. (2026).

Building an MLX backend

To bring K-Search to Apple Silicon, we first built a native MLX backend. We implemented a full MLX-specific task adapter for K-Search, including:

• An MLX task backend in k_search/tasks/ handling kernel compilation and execution on Apple Silicon via MLX’s Metal/C++ APIs.
• Updated kernel generator prompts for writing and modifying Metal/MLX kernels.
• MLX-specific benchmarking integration using mlx.core measurement utilities.

Translating CUDA expertise to MLX

However, the more interesting challenge was not simply running K-Search on MLX. The key insight is that expert CUDA kernels encode decades of optimization knowledge that is transferable to Apple GPU if you can bridge the conceptual gap. Simply handing an LLM a CUDA kernel and asking it to port it is not enough: without deep hardware context, it produces code that is syntactically valid but architecturally wrong (wrong tile sizes, invalid primitives, mismatched memory assumptions).

Our translation layer consists of:

• Concept mapping tables: A structured glossary of CUDA primitives and their MLX/Metal equivalents with hard constraints. For example:
  • __shared__ maps to Metal threadgroup memory but with a hard 32 KB limit (vs. NVIDIA’s 48 KB)
  • warp_reduce maps to MMA (preferred)
  • __syncthreads() becomes threadgroup_barrier(mem_flags::mem_tg)
  • H100’s ~3.35 TB/s HBM3 maps to M3 Max’s ~400 GB/s unified DRAM a bandwidth difference that reshapes which optimizations are worth pursuing.
• MLX-specific hints and patterns: Concrete code-level patterns for operations with no direct CUDA equivalent, such as register-based row reductions using simd_shuffle_xor in an 8×8 MMA tile layout, or the “exp2 trick” (replacing $exp(x)$ with $exp_2(x \log_2 e)$) for faster softmax on Apple’s fast $exp_2$ hardware instruction.
• Reusable assertions: Expert kernel behaviors reframed as properties the evolutionary search must preserve, rather than code to copy.

Matching expert kernel performance: the Attention kernel

We evaluate three configurations of an MLX attention kernel for Apple Silicon: (1) a naive baseline, (2) pure evolution with no additional provided context, and (3) a full context translation layer, which supplies the optimizer with architecture-specific implementation knowledge extracted from high-performance kernels (e.g., FlashAttention-2), letting the evolutionary search reason about implementation strategies rather than starting from a naive kernel. Together, these three configurations let us isolate the exact impact of the translation layer.

Figure 4: Performance scaling of the Attention Kernel through stacked optimizations. The "Full Context" configuration successfully discovers and implements advanced strategies like double buffering and loop unrolling, achieving near-expert performance.

The jump from 0.26× to 0.97× the speed of Apple’s state-of-the-art attention kernel — illustrates how much the translation layer matters. With full context, the evolved kernel independently discovers the key optimizations in FlashAttention 2: threadgroup memory tiling, online softmax, K-transposition for memory access, and the exp2 trick. The last of these replaces every softmax exponential with a base-2 exponential,

\[e^x = 2^{x \log_2 e},\]

which is exact and lets the kernel use Apple’s fast fast::exp2() hardware instruction directly instead of paying for a base conversion at runtime.

A 20× faster prefill: the Mamba SSM kernel

To evaluate whether K-Search generalizes beyond attention kernels, we applied it to the state-space model (SSM) kernel used by Mamba. Unlike attention, the computational bottleneck is a recurrent state update rather than a softmax, providing a substantially different optimization challenge. We compare the evolved implementation against the community MLX implementation (mlx-lm) and the PyTorch reference implementation (mamba.py) on an M1 Max.

Evaluated on mamba-370m f16, M1 Max 64GB:

Metric	mlx-mamba (ours)	mlx-lm (community)	mamba.py
Decode	152 tok/s	116 tok/s	40 tok/s
Prefill L=512	5,751 tok/s	329 tok/s	1,089 tok/s
Prefill L=1024	6,010 tok/s	327 tok/s	1,127 tok/s
Prefill L=2048	6,612 tok/s	326 tok/s	1,092 tok/s
Prefill L=4096	6,743 tok/s	339 tok/s	1,042 tok/s

Table 1: Prefill and decode throughput on mamba-370m (f16, M1 Max 64GB). mlx-mamba (ours) reaches ~20× higher prefill throughput than the community mlx-lm baseline, while decode remains comparable.

The ~20× prefill speedup over mlx-lm comes down to one difference: mlx-lm does not implement a parallel scan for the SSM. The state recurrence

\[h_t = \bar{a}t h{t-1} + \bar{b}_t\]

looks inherently sequential, but each step can be written as a pair $(\bar{a}_t, \bar{b}_t)$ under the associative combine

\[(a_2, b_2) \circ (a_1, b_1) = \left(a_2 a_1,\ a_2 b_1 + b_2\right),\]

which reproduces the recurrence exactly. Because the operator is associative, the whole sequence can be evaluated with a parallel (prefix) scan in $O(\log N)$ dependent steps instead of $O(N)$. mlx-lm skips this and processes tokens one at a time, leaving most of Apple Silicon’s compute idle; our evolved Metal kernel applies the scan and makes much fuller use of GPU throughput. The gain shows up in prefill, where the full sequence is available to scan in parallel, and not in single-token decode, where there is only one new token per step and no scan to parallelize — which is why the decode row is roughly flat while prefill is ~20×.

mamba.py is slow on both prefill and decode because it is a PyTorch reference implementation that falls back to CPU or MPS on Apple Silicon, forgoing the hardware-specific optimizations that MLX’s Metal backend makes possible.

What’s next?

On the two kernels we studied, AI-driven evolutionary kernel search grounded in structured cross-platform translation knowledge reached near-expert performance on Apple Silicon without a team of GPU experts starting from scratch. We do not yet know how far this generalizes, but the result is encouraging.

For us the main takeaway is that the bottleneck was not the LLM’s ability to write Metal code, but the quality of the context and constraints we gave it. Our CUDA translation layer converts existing NVIDIA kernel expertise into actionable guidance for Apple Silicon, and lets K-Search’s evolutionary search do the rest.

We are actively extending this work in several directions: supporting new architectures, with current efforts focused on developing new kernels for the IBM Spyre AIU and broader hardware targets; adding more kernels such as paged attention and fused MoE routing; and improving integration with the K-Search evolution loop to make translation context even more automatic.

Acknowledgements

This work was carried out by IBM Research and builds on K-Search from the UC Berkeley Sky Lab (Cao et al., 2026). We welcome collaboration and feedback from the MLX and broader AI systems communities. If you are working on kernel optimization for non-CUDA hardware, we would love to hear from you.

---

Citation

@article{cao2026k,
  title={K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model},
  author={Cao, Shiyi and Mao, Ziming and Gonzalez, Joseph E and Stoica, Ion},
  journal={arXiv preprint arXiv:2602.19128},
  year={2026}
}

Appendix: Try it yourself

The MLX backend is built on top of the open-source K-Search repo, so the results here can be reproduced directly. The steps are:

1. Clone and install

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

2. Set your credentials

Open the relevant script under scripts/ and set three variables at the top:

KSEARCH_ROOT=/path/to/K-Search
API_KEY=your-llm-api-key

3. Run kernel search

# Optimize Flash Attention on Apple Silicon (world-model mode)
bash scripts/mac_flash_attention_wm.sh

# Or a Mamba SSM kernel, e.g. the selective scan
bash scripts/mamba_selective_scan_fwd_wm.sh

Full CLI reference and documentation are in the README.