跳到正文
原文
Berkeley AI Research·· 2026-07-29精选AI 评分60

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

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

AI 导读

Berkeley Sky Lab 与 IBM Research 将 K-Search 这一进化式内核搜索框架扩展到 MLX，用结构化的 CUDA-to-MLX 翻译层，把已有 CUDA 内核当作知识库迁移到 Apple Silicon。

推荐理由

这篇工作给出把 CUDA 内核经验迁移到 Apple Silicon 的具体做法，读者可据此判断跨平台内核优化的可行路径。

正文 · AI 翻译

Kernel knowledge transfer from CUDA to MLX

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

我们正面临计算领域的一个新时代。硬件在快速变化——不仅仅是更快的 GPU，还有来自不同厂商的日益增多的芯片，每种芯片都有自己的架构，并且往往针对特定的 AI 工作负载进行了定制。软件同样在快速变化，AI 编程工具现在能在几分钟内生成几年前需要数月努力才能完成的东西。

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

在这项工作中，我们探讨这种专业知识能否被自动转移。我们基于 K-Search 构建——这是 Cao 等人在 Berkeley Sky Lab 提出的进化式内核搜索框架，它利用 AI 优化 GPU 内核——并为其扩展了一个用于 MLX 的后端，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 推理接下来应尝试哪些优化，一个代码编写模型生成候选内核，然后这些候选内核在真实硬件上被编译和基准测试。

测量结果反馈到搜索中，搜索不断精炼，追踪有前景的方向并淘汰死胡同，直到性能收敛。

Pseudocode for K-Search via co-evolving world models

算法 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 世界模型节点示例。 每个候选优化记录一个具体动作、估计的硬件影响、一个总体优先级评分以及模型的置信度。

Overview of the K-Search loop

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

原始 K-Search 论文在 FlashInfer 的 CUDA 内核上评估了这一搜索策略。在 GQA decode、MLA decode、MLA prefill 和 MoE 上，在相同的 120 次迭代预算下，K-Search 相比 OpenEvolve 和 ShinkaEvolve 提升得更稳定。这些结果奠定了我们在此构建的搜索框架；本文的剩余部分要问的是，它的优化知识能否迁移到 CUDA 之外。

K-Search benchmark results compared with OpenEvolve and ShinkaEvolve

图 3：原始 K-Search 论文的主要结果。在三次运行中，K-Search 在四个 FlashInfer CUDA 内核上均取得了比 OpenEvolve 和 ShinkaEvolve 更强的 best-so-far 搜索得分、单工作负载内核性能和加速比分布。完全复制自 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 进行基于寄存器的行归约，或“exp2 技巧”（将 $exp(x)$ 替换为 $exp_2(x \log_2 e)$），以在 Apple 快速的 $exp_2$ 硬件指令上实现更快的 softmax。
  • 可复用的断言：将专家级内核行为重新表述为演化搜索必须保留的属性，而不是要复制的代码。

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

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

Performance scaling of the Attention Kernel through stacked optimizations

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

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

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

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

预填充快 20×：Mamba SSM 内核

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

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

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

表 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，让苹果芯片的大部分算力闲置；我们进化出的 Metal 内核应用了扫描，更充分地利用了 GPU 吞吐量。这一收益体现在预填充中，因为此时整个序列可并行扫描；而在单 token 解码中则不然，因为每步只有一个新 token，没有可并行的扫描——这就是为什么解码一行大致持平，而预填充约为 20×。

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

下一步是什么？

在我们研究的这两个内核上，由结构化跨平台翻译知识驱动的 AI 进化内核搜索，在苹果芯片上达到了接近专家的性能，而无需一支 GPU 专家团队从零开始。我们尚不知道这能推广到多远，但结果令人鼓舞。

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

我们正在多个方向上积极推进这项工作：支持新架构，目前的工作重点是为 IBM Spyre AIU 及更广泛的硬件目标开发新内核；添加更多内核，例如 paged attention 和融合 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。

来源：Berkeley AI Research · bair.berkeley.edu

#论文/研究#推理#部署/工程#开源生态