从 CUDA 到 MLX：K-Search 如何把内核优化经验带到 Apple Silicon

图 1：CUDA 到 MLX 优化映射图。CUDA 优化知识可以转化为架构原生的 MLX 策略，而不是逐条指令照搬。

我们正面临计算领域的新纪元。硬件正在快速变化——不仅是更快的 GPU，还有来自不同厂商、种类日益增多的芯片，每一种都有其自身的架构，且往往针对特定的 AI 工作负载量身定制。软件的变化同样迅速，AI 编程工具如今几分钟就能生成几年前需要数月努力才能完成的东西。

如今计算的重心如此集中于 AI，GPU 内核是其成功的关键组成部分。这些是在 GPU 内部运行的低层程序，而编写高效的内核绝非易事——需要多年的专业经验才能做好。将内核从一个厂商的硬件迁移到另一个厂商的硬件则更加困难，往往意味着要从头重新发现同样的优化。例如，CUDA 生态系统积累了数十年来之不易的内核专业知识：手工调优的注意力、状态空间模型以及其他关键操作的实现，代表着数千小时的工程投入。较新的硬件生态系统（Apple Silicon、定制 AI 加速器等）发展迅速，但缺乏这种深度。

在这项工作中，我们探究这种专业知识能否被自动迁移。我们基于 K-Search 展开——这是 Cao 等人在 Berkeley Sky Lab 提出的进化式内核搜索框架，利用 AI 优化 GPU 内核——并为其扩展了一个面向 MLX（Apple 为其自家 Apple Silicon 芯片打造的机器学习框架）的后端。我们开发了一种新颖的结构化 CUDA 到 MLX 翻译层，使 K-Search 能够以现有 CUDA 内核作为知识库，将其适配为适用于 Apple Silicon 的高质量 GPU 内核，而非从头重建。

我们展示了该方法在 Apple Silicon 上达到了接近专家级的性能，与原生 MLX Attention 内核相比实现了 0.97 倍的加速，在 Mamba SSM 内核上相比社区 mlx-lm 实现最高实现了 20 倍的预填充加速；我们在下文各节中报告了这些数字，以及其中有多少收益来自翻译层。尽管我们聚焦于面向 Apple Silicon 的 MLX 内核，但该方法并不局限于 MLX，适用于任何 CUDA 专业知识可迁移的生态系统。

为什么选择 MLX？

自 2023 年底以来，Apple 的 MLX 框架获得了显著的应用。随着 Apple Silicon 出现在数亿台 MacBook 和 Mac Studio 中，MLX 使得本地 AI 推理无需云成本即可实现。统一内存架构使其对中型模型（M 系列芯片上 7B–70B 参数）尤为有吸引力。

然而，在这一势头之下存在一个显著缺口：NVIDIA 生态系统视为理所当然的许多性能关键型内核：分页注意力、优化的 SSM 扫描内核、融合 MoE 路由——要么缺失，要么在没有硬件特定调优的情况下显得简陋。MLX 能正确运行模型，但往往在性能上留下了大量未挖掘的潜力。

这一缺口正是本文后续内容的动机所在。

什么是 K-Search？

K-Search 是一个进化式内核优化框架，最初由我们的第一作者 Shiyi Cao 在 UC Berkeley Sky Lab 开发。给定一个朴素内核和一份硬件规格，它会运行一个迭代优化循环：一个 LLM 推理接下来该尝试哪些优化，一个代码编写模型生成候选内核，然后这些候选内核在真实硬件上被编译和基准测试。

测量结果反馈到搜索中，搜索不断细化，推进有前景的方向，舍弃死胡同，直到性能收敛。

算法 1：通过协同演化世界模型的 K-Search。 搜索在以下步骤之间交替：选择最有前景的动作，实例化并评估代码直到改进停滞，以及通过插入、更新和剪枝操作演化世界模型。改编自 Cao et al. (2026)。

搜索以 Spec 为基础：一份领域特定文档，编码了硬件规则、优化模式和数学约束，防止生成的代码幻觉出无效原语，并确保候选方案能够真正编译并高效运行。

在我们的运行中，单个模型（Gemini 3.5 Pro Preview）同时扮演两个角色：它维护推理状态并编写内核。推理部分被提示为“GPU 内核性能工程师”，并被要求在提出任何建议之前完成固定分析：对内核进行分类（归约、扫描、注意力/softmax……），将参考计算重写为规范形式，梳理数据布局和访问模式，并假设每种运行时状态下可能的瓶颈（带宽、延迟、计算或同步）。只有在这之后，它才会输出候选优化，每个优化都是可在一次迭代中实现的单一更改。

我们将这种持久推理状态称为世界模型。它不是一份扁平的待尝试事项列表，而是一棵决策（前缀）树：每条根→叶路径组合成一个完整的优化计划，兄弟分支是相互竞争的备选方案。每个节点都会被评分——一个 [0, 10] 范围内的 overall_rating、一个 [0, 1] 范围内的 confidence，以及每个节点在内存带宽、寄存器压力和计算/硬件匹配度上的 impacts——这样搜索就可以对部分计划进行排序，并扩展最有前景的计划。这棵树会跨轮次持续存在并生长：细化一个想法会添加一个子节点，而不是覆盖其父节点；如果最佳分数在几轮内未能提高（停滞窗口），搜索就会回退以探索另一个分支。单个节点在注意力内核运行中途的样子如下：

{
  "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 节点前沿（橙色，待定假设如 $u_{13}$）组成。工作流迭代三个阶段的循环：(1) 动作选择，根据世界模型估计的优先级分数 $V$ 从前沿中检索最有前景的动作节点；(2) 局部细化，其中随机策略 $\pi_{\mathrm{code}}$ 采样具体实现直到停滞；以及 (3) 世界模型更新，其中 LLM 对轨迹进行推理，通过插入（添加新动作）、更新（调整 $V$，例如 $u_{11}$ 从 0.9 降至 0.6）和剪枝（移除如 $u_{10}$ 这样前景较差的节点）来更新搜索树。

最初的 K-Search 论文在 FlashInfer 的 CUDA 内核上评估了这一搜索策略。在 GQA 解码、MLA 解码、MLA 预填充和 MoE 上，在相同的 120 次迭代预算下，K-Search 的改进比 OpenEvolve 和 ShinkaEvolve 更稳定。这些结果确立了我们在本文中所基于的搜索框架；本文的其余部分则探讨其优化知识能否迁移到 CUDA 之外。

图 3：原始 K-Search 论文的主要结果。在三次运行中，K-Search 在四个 FlashInfer CUDA 内核上取得了比 OpenEvolve 和 ShinkaEvolve 更强的历史最佳搜索分数、各工作负载内核性能和加速比分布。完全复现自 Cao et al. (2026)。

构建 MLX 后端

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

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

将 CUDA 专业知识迁移到 MLX

然而，更有趣的挑战并非仅仅是在 MLX 上运行 K-Search。关键洞察在于，专家级 CUDA 内核编码了数十年的优化知识，只要你能弥合概念上的差距，这些知识就可以迁移到 Apple GPU。仅仅把一个 CUDA 内核交给 LLM 并要求它移植是不够的：如果没有深入的硬件背景，它生成的代码虽然语法有效，但在架构上是错误的（错误的 tile 大小、无效的原语、不匹配的内存假设）。

我们的翻译层包括：

• Concept mapping tables: A structured glossary of CUDA primitives and their MLX/Metal equivalents with hard constraints. For example:
  • __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 进行基于寄存器的行归约，或利用 Apple 快速 $exp_2$ 硬件指令的“exp2 技巧”（将 $exp(x)$ 替换为 $exp_2(x \log_2 e)$）以加速 softmax。
• 可复用断言：将专家内核行为重新表述为进化搜索必须保持的属性，而非要复制的代码。

匹配专家内核性能：Attention 内核

我们评估了面向 Apple Silicon 的 MLX attention 内核的三种配置：(1) 朴素基线，(2) 不提供额外上下文的纯进化，(3) 完整上下文翻译层，它为优化器提供从高性能内核（如 FlashAttention-2）中提取的架构专用实现知识，让进化搜索能够推理实现策略，而不是从朴素内核开始。这三种配置共同让我们能够分离出翻译层的确切影响。

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

从 0.26× 提升到 Apple 最先进注意力内核速度的 0.97×——这一跃升说明了翻译层有多么重要。在完整上下文中，进化出的内核独立发现了 FlashAttention 2 中的关键优化：线程组内存分块、在线 softmax、用于内存访问的 K 转置，以及 exp2 技巧。其中最后一项将每个 softmax 指数替换为以 2 为底的指数，

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

这是精确的，并让内核可以直接使用 Apple 快速的 fast::exp2() 硬件指令，而不必在运行时付出底数转换的代价。

预填充速度提升 20×：Mamba SSM 内核

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

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

表 1：mamba-370m（f16，M1 Max 64GB）上的预填充和解码吞吐量。mlx-mamba（我们的）预填充吞吐量比社区 mlx-lm 基线高约 20×，而解码保持相当。

相比 mlx-lm 约 20× 的预填充加速归结为一个差异：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)$ 个依赖步骤内求值，而不是 $O(N)$。mlx-lm 跳过了这一点，逐个处理 token，使 Apple Silicon 的大部分算力闲置；我们进化出的 Metal 内核应用了扫描，更充分地利用了 GPU 吞吐量。这一收益体现在预填充中，因为此时整个序列可用于并行扫描，而不体现在单 token 解码中，因为每步只有一个新 token，没有可并行化的扫描——这就是为什么解码行大致持平，而预填充约为 20×。

mamba.py 在预填充和解码上都很慢，因为它是一个 PyTorch 参考实现，在 Apple Silicon 上回退到 CPU 或 MPS，放弃了 MLX 的 Metal 后端所能实现的硬件特定优化。

下一步是什么？

在我们研究的两个内核上，以结构化跨平台翻译知识为基础、由 AI 驱动的进化内核搜索，在 Apple Silicon 上达到了接近专家的性能，而无需一个从零开始的 GPU 专家团队。我们尚不知道这能泛化到何种程度，但结果令人鼓舞。

对我们来说，主要结论是瓶颈不在于 LLM 编写 Metal 代码的能力，而在于我们提供给它的上下文和约束的质量。我们的 CUDA 翻译层将现有的 NVIDIA 内核专业知识转化为针对 Apple Silicon 的可操作指导，并让 K-Search 的进化搜索完成其余工作。

我们正在多个方向上积极扩展这项工作：支持新架构，目前的工作重点是针对 IBM Spyre AIU 和更广泛的硬件目标开发新内核；添加更多内核，例如分页注意力和融合 MoE 路由；以及改进与 K-Search 演化循环的集成，使翻译上下文更加自动化。

致谢

这项工作由 IBM Research 完成，并建立在 UC Berkeley Sky Lab 的 K-Search（Cao et al., 2026）之上。我们欢迎来自 MLX 及更广泛 AI 系统社区的合作与反馈。如果你正在从事非 CUDA 硬件的内核优化工作，我们很乐意听到你的声音。

引用

@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}
}

附录：自己动手试试

MLX 后端构建在开源 K-Search 仓库之上，因此这里的结果可以直接复现。步骤如下：

1. 克隆并安装

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. 设置你的凭据

打开 scripts/ 下的相关脚本，并在顶部设置三个变量：

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

3. 运行内核搜索

# 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

完整的 CLI 参考和文档见 README。