CUDA 优化迁移到 Apple Silicon：内核知识翻译指南

CUDA 优化迁移到 Apple Silicon：内核知识翻译指南 | 前端进阶之旅

Full translation

完整中文译文

序言

我们正处于计算的新纪元。硬件变化迅速——不仅仅是 GPU 更快，还有来自不同厂商的芯片日益丰富，每种都有自己的架构，常常针对特定的 AI 工作负载进行优化。软件变化同样迅速，AI 编码工具现在能在几分钟内生成数年前需要花费数月才能完成的代码。

随着计算日益围绕 AI 展开，GPU 内核成为其成功的关键组件。这些是在 GPU 内部运行的低级程序，编写高效的内核远非易事——需要多年的专业积累才能做好。将一个内核从一个厂商的硬件迁移到另一个，难度更大，通常意味着要从零开始重新发现相同的优化。例如，CUDA 生态系统已经积累了数十年的内核优化专业知识：手工调优的 attention、state space models 和其他关键操作的实现，代表了数千个工程小时的投入。较新的硬件生态系统（Apple Silicon、定制 AI 加速器等）增长迅速，但在这方面深度不足。

在这项工作中，我们探索这些专业知识是否能够自动迁移。我们基于 K-Search——由 Berkeley Sky Lab 的 Cao 等人引入的进化型内核搜索框架，该框架使用 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 倍的 prefill 加速；我们在下面的章节中报告这些数字和多少收益来自转换层。虽然我们关注的是 Apple Silicon 上的 MLX 内核，但该方法并非特定于 MLX，适用于任何 CUDA 专业知识可转移的生态系统。

背景：Apple Silicon 上的性能缺口

Apple 的 MLX 框架自 2023 年末以来获得了显著的应用。凭借 Apple Silicon 在数亿部 MacBook 和 Mac Studio 中的存在，MLX 使本地 AI 推理无需云成本。统一内存架构使其对中等规模模型（M 系列芯片上的 7B–70B 参数）特别有吸引力。

然而，这种势头下存在一个重大缺口：NVIDIA 生态系统司空见惯的许多性能关键内核——分页 attention、优化的 SSM 扫描内核、融合 MoE 路由——要么缺失，要么没有硬件特定的调优。MLX 能够正确运行模型，但往往在性能上还有很大空间。

这个缺口正是本文其余部分的动力。

K-Search：进化型内核优化框架

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

测量结果反馈到搜索中，搜索不断优化，追求有前景的方向并放弃死胡同，直到性能收敛。

Algorithm 1: K-Search via co-evolving world models. 搜索在选择最有前景的动作、实例化和评估代码直到改进停滞、以及通过 insert、update 和 prune 操作演进世界模型之间交替进行。改编自 Cao et al. (2026)。

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

在我们的运行中，单一模型（Gemini 3.5 Pro Preview）扮演两个角色：它维护推理状态并编写内核。推理部分被提示为一个"GPU 内核性能工程师"，被要求在提议任何内容之前进行固定的分析：分类内核（reduction、scan、attention/softmax 等），以规范形式重写参考计算，绘制数据布局和访问模式，并假设每个运行时制度中可能的瓶颈（带宽、延迟、计算或同步）。只有这样才会发出候选优化，每个都作为单一改变，可在一次迭代中实现。

我们称持久推理状态为世界模型。它不是一个扁平的列表，而是一个决策（前缀）树：每条从根到叶的路径组成一个完整的优化计划，兄弟分支是竞争的替代方案。每个节点都被评分——overall_rating 在 [0, 10]、confidence 在 [0, 1]，以及每个节点对内存带宽、寄存器压力和计算/硬件适配的影响——因此搜索可以排列部分计划，并扩展最有前景的。树在各轮之间持久化并增长：优化一个想法添加一个子节点而不是覆盖其父节点，如果最佳评分在几轮内未能改进（停滞窗口），搜索会退回来探索替代分支。单个节点在运行中的 attention 内核上如下所示：

{
  "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. 每个候选优化记录一个具体的动作、估计的硬件影响、总体优先级评分和模型的置信度。

Figure 2: Overview of K-Search. 框架在搜索状态 $S_t$ 上运行，结构化为搜索树。树由已访问节点（蓝色，已访问的状态，附加程序如 $x_{12}$）和前沿的开放节点（橙色，待处理的假设如 $u_{13}$）组成。工作流通过三个阶段迭代：(1) 动作选择，其中从前沿检索最有前景的动作节点，基于世界模型估计的优先级评分 $V$；(2) 局部优化，其中随机策略 $\pi_{\mathrm{code}}$ 采样具体实现直到停滞；以及 (3) 世界模型更新，其中 LLM 推理过轨迹以通过 Insert（添加新动作）、Update（调整 $V$，如 $u_{11}$ 从 0.9 降至 0.6）和 Prune（移除较少有前景的节点如 $u_{10}$）来更新搜索树。

原始 K-Search 论文在 FlashInfer 中的 CUDA 内核上评估了这个搜索策略。在 GQA decode、MLA decode、MLA prefill 和 MoE 上，K-Search 在相同 120 次迭代预算下的改进比 OpenEvolve 和 ShinkaEvolve 更一致。这些结果确立了我们基于的搜索框架；本文其余部分询问其优化知识是否能超越 CUDA 迁移。

Figure 3: Main results from the original K-Search paper. 在三次运行中，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 并要求其移植是不够的：没有深入的硬件上下文，它会生成语法上有效但架构上错误的代码（错误的瓦片大小、无效的原语、不匹配的内存假设）。

我们的转换层包括：

概念映射表： 一个 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 瓦片布局中使用 simd_shuffle_xor 的基于寄存器的行约简，或"exp2 trick"（将 $exp(x)$ 替换为 $exp_2(x \log_2 e)$），以在 Apple 快速 $exp_2$ 硬件指令上加速 softmax。

可复用的断言： 专家内核行为重新表述为进化搜索必须保留的属性，而不是要复制的代码。

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

我们评估了 Apple Silicon 上 MLX attention 内核的三个配置：(1) 一个朴素基线，(2) 纯进化，无额外提供的上下文，以及 (3) 完整上下文转换层，为优化器提供从高性能内核（如 FlashAttention-2）提取的架构特定实现知识，让进化搜索能够推理实现策略，而不是从朴素内核开始。这三个配置共同让我们隔离转换层的精确影响。

Figure 4: Performance scaling of the Attention Kernel through stacked optimizations. "Full Context" 配置成功发现并实现了如双缓冲和循环展开的高级策略，达到了接近专家的性能。

从 0.26× 到 0.97× Apple 最先进 attention 内核速度的跳跃——说明转换层有多重要。有完整上下文的情况下，进化内核独立发现了 FlashAttention 2 中的关键优化：threadgroup 内存瓦片、在线 softmax、K 转置以改进内存访问，以及 exp2 trick。最后一个将每个 softmax 指数替换为以 2 为底的指数，

这是精确的，让内核直接使用 Apple 的快速 fast::exp2() 硬件指令，而不是在运行时支付基数转换的成本。

20 倍更快的 prefill：Mamba SSM 内核

为了评估 K-Search 是否泛化超越 attention 内核，我们将其应用于 Mamba 使用的 state-space model (SSM) 内核。与 attention 不同，计算瓶颈是循环状态更新而不是 softmax，提供了实质不同的优化挑战。我们在 M1 Max 上将进化实现与社区 MLX 实现（mlx-lm）和 PyTorch 参考实现（mamba.py）进行比较。

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

Table 1: Prefill and decode throughput on mamba-370m (f16, M1 Max 64GB). mlx-mamba（我们的）相比社区 mlx-lm 基线达到 ~20× 更高的 prefill 吞吐量，而 decode 保持可比。

相比 mlx-lm 的 ~20× prefill 加速归结为一个差异：mlx-lm 未实现 SSM 的并行扫描。状态循环

看起来固有的顺序，但每一步都可以写成一对 $(\bar{a}_t, \bar{b}_t)$ 在关联组合

精确再现循环。因为算子是关联的，整个序列可以用并行（前缀）扫描在 $O(\log N)$ 依赖步而不是 $O(N)$ 内评估。mlx-lm 跳过这个，一次处理一个 token，使 Apple Silicon 的大多数计算闲置；我们进化的 Metal 内核应用扫描，更充分地利用 GPU 吞吐量。收益显示在 prefill 中，整个序列可用于并行扫描，而不是在单 token decode 中，其中每步只有一个新 token，没有扫描可并行化——这就是为什么 decode 行基本平坦而 prefill 是 ~20×。

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

结论

在我们研究的两个内核上，基于结构化跨平台转换知识的 AI 驱动进化内核搜索在 Apple Silicon 上达到了接近专家的性能，无需从零开始的 GPU 专家团队。我们还不知道这泛化到多远，但结果令人鼓舞。

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

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

这项工作由 IBM Research 开展，基于来自 UC Berkeley Sky Lab（Cao et al., 2026）的 K-Search。我们欢迎来自 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 repo 之上，所以这里的结果可以直接再现。步骤是：

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 中。