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

从 CUDA 到 MLX：K-Search 如何将数十年内核经验带入 Apple 芯片 – 伯克利人工智能研究博客

twitter

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

我们正面临计算领域的新纪元。硬件正在快速演变——不仅仅是更快的 GPU，而是来自不同厂商的芯片种类日益增多，每种芯片都有其独特的架构，通常针对特定的 AI 工作负载进行优化。软件的演变同样迅速，如今 AI 编码工具只需几分钟即可完成过去需要数月努力才能实现的功能。

more

随着计算越来越多地围绕 AI 展开，GPU 内核已成为其成功的关键组成部分。这些是运行在 GPU 内部的底层程序，编写高效的内核远非显而易见——需要数年的专业知识才能掌握。将一个内核从一种硬件转移到另一种硬件更加困难，通常意味着需要从头开始重新发现相同的优化方法。例如，CUDA 生态系统已经积累了数十年来通过实践获得的内核专业知识：包括注意力机制、状态空间模型和其他关键操作的手工优化实现，这些代表了数千小时的工程工作。而新的硬件生态系统（如 Apple 芯片、定制 AI 加速器等）虽然发展迅速，但缺乏这种深度。

在本研究中，我们探讨是否可以自动转移这些专业知识。我们基于 Cao 等人在伯克利 Sky Lab 提出的 K-Search 演化内核搜索框架（该框架使用 AI 优化 GPU 内核），并为其扩展了 MLX 后端——苹果公司为其 Apple 芯片设计的机器学习框架。我们开发了一种新颖的结构化 CUDA 到 MLX 翻译层，使 K-Search 能够将现有的 CUDA 内核作为知识库，将其适配为 Apple 芯片的高质量 GPU 内核，而非从零开始重新构建。

我们证明，我们的方法在 Apple 芯片上实现了接近专家水平的性能，相比原生 MLX 注意力内核，速度提升了 0.97 倍；在 Mamba SSM 内核上，相比社区实现的 mlx-lm，预填充速度最高提升了 20 倍；以下部分将报告具体数值，并说明这些提升中有多少来自于翻译层。尽管我们专注于 Apple 芯片的 MLX 内核，但该方法并不仅限于 MLX，适用于任何可以转移 CUDA 专业知识的生态系统。

为什么是 MLX？

自 2023 年底以来，苹果的 MLX 框架获得了显著的采用。凭借搭载 Apple 芯片的数亿台 MacBook 和 Mac Studio，MLX 使得本地 AI 推理无需依赖云端计算，从而节省了云服务成本。统一内存架构使其在中型模型（M 系列芯片上的 7B-70B 参数）上尤为具有吸引力。

然而，在这种发展势头之下，存在一个显著的差距：NVIDIA 生态系统中许多关键的性能优化内核（如分页注意力、优化后的 SSM 扫描内核、融合的 MoE 路由等）在未进行硬件特定调优时要么缺失，要么表现原始。MLX 虽然能正确运行模型，但往往未能充分发挥性能潜力。

正是这一差距，促使我们继续探讨后续内容。

什么是 K-Search？

K-Search 是一个进化型内核优化框架，最初由本文第一作者曹世毅在加州大学伯克利分校 Sky Lab 开发。给定一个原始内核和硬件规格，它会运行一个迭代优化循环：大语言模型（LLM）推理下一步应尝试哪些优化，代码生成模型生成候选内核，这些候选方案会在真实硬件上进行编译和基准测试。

测量结果会反馈到搜索过程中，持续优化方向，探索有前景的路径并舍弃死胡同，直到性能趋于收敛。

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

搜索过程由一个 Spec 作为基础：这是一个领域特定文档，编码了硬件规则、优化模式和数学约束，防止生成代码出现无效原语的幻觉，并确保候选方案能够实际编译运行且高效。

在我们的实验中，单个模型（Gemini 3.5 Pro Preview）同时承担两个角色：维护推理状态并编写内核。推理部分被提示为"GPU 内核性能工程师"，要求在提出任何建议前完成固定分析流程：对内核进行分类（归约、扫描、注意力/softmax 等），将参考计算重写为规范形式，绘制数据布局和访问模式，并在每个运行时场景中推测可能的瓶颈（带宽、延迟、计算或同步）。只有在此之后，才会生成候选优化方案，每个方案都作为单次迭代可实现的单一修改。

我们将这种持续的推理状态称为世界模型。它不是一个简单的待尝试事项列表，而是一个决策（前缀）树：每个从根到叶的路径构成完整的优化方案，兄弟分支则是竞争性替代方案。每个节点都有评分——总体评分（0-10）、置信度（0-1），以及对内存带宽、寄存器压力和计算/硬件适配的节点级影响——使搜索能够对部分方案进行排序并扩展最有前景的方案。该树结构在多轮搜索中持续存在并增长：改进想法通过添加子节点进行细化而非覆盖父节点，如果最佳评分在若干轮（停滞窗口）内未改善，搜索会回退以探索替代分支。在注意力内核运行中途，单个节点的示例如下：

{
"action"
:
"将线程组内存 softmax 归约替换为仅寄存器归约：每个 SIMD 组拥有 8 个查询行，并通过 simd_shuffle_xor 在 lane 间归约，移除线程组屏障。",
"difficulty_1_to_5"
:
4
,
"impacts"
:
{
"memory_bandwidth"
:
8
,
"register_pressure"
:
4
,
//
风险：如果 Br > 8 可能发生溢出
"compute_hw_fit"
:
9
//
SIMD 宽度 32；保持 8x8 分块
},
"overall_rating_0_to_10"
:
8
,
"confidence_0_to_1"
:
0.7
}

列表 1：K-Search 世界模型节点示例。每个候选优化记录具体操作、预估硬件影响、总体优先级评分以及模型的置信度。

Figure 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 任务中，K-Search 在相同 120 次迭代预算下比 OpenEvolve 和 ShinkaEvolve 更一致地提升了性能。这些结果确立了我们在此基础上构建的搜索框架；本文其余部分将探讨其优化知识是否能超越 CUDA 领域。

Figure 3: 原始 K-Search 论文的主要结果。在三次运行中，K-Search 在四个 FlashInfer CUDA 内核上实现了比 OpenEvolve 和 ShinkaEvolve 更强的最优搜索得分、每工作负载内核性能和加速分布。完全复现自 Cao 等人（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

然而，更具挑战性的并非仅仅是让 K-Search 在 MLX 上运行。关键洞察在于：专家级 CUDA 内核编码了数十年的优化知识，若能弥合概念鸿沟，这些知识可迁移至 Apple GPU。简单地将 CUDA 内核交给 LLM 并要求其移植是不够的：缺乏深度硬件背景时，生成的代码虽语法正确但架构错误（错误的分块大小、无效原语、内存假设不匹配）。

我们的翻译层包含：

• 概念映射表：结构化的 CUDA 原语与 MLX/Metal 等价项对照表，包含硬性约束。例如：__shared__ 映射到 Metal 线程组内存但有 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 实现基于寄存器的行归约，或利用 Apple 快速 exp_2 硬件指令的 "exp2 技巧"（将 $exp(x)$ 替换为 $exp_2(x \log_2 e)$）加速 softmax 计算。

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

与专家内核性能匹配：注意力内核

我们评估了用于 Apple Silicon 的 MLX 注意力内核的三种配置：(1) 一个原始基线，(2) 无额外上下文的纯进化方法，(3) 完整上下文翻译层，该层向优化器提供从高性能内核（如 FlashAttention-2）中提取的架构特定实现知识，使进化搜索能够推理实现策略，而非从原始内核开始。这三种配置共同使我们能够隔离翻译层的确切影响。

图 4：通过分层优化实现的注意力内核性能扩展。"完整上下文"配置成功发现了并实现了双重缓冲和循环展开等高级策略，达到接近专家水平的性能。

从 0.26× 到 0.97× Apple 最先进注意力内核速度的飞跃——展示了翻译层的重要性。在完整上下文中，进化后的内核独立发现了 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 上评估：

| 指标 | 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 跳过了这一步骤，逐个处理标记，导致 Apple Silicon 的大部分计算资源处于空闲状态；我们改进的 Metal 内核则应用了扫描，更充分地利用了 GPU 的吞吐能力。这一优势体现在预填充阶段，此时整个序列可用于并行扫描，而在单标记解码阶段则没有可并行化的扫描操作（每个步骤只有一个新标记），因此解码行的性能基本保持平稳，而预填充性能约为 20 倍。

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

下一步计划

在我们研究的两个内核上，基于结构化跨平台翻译知识的 AI 驱动进化内核搜索方法，在无需 GPU 专家团队从零开始的情况下，就在 Apple Silicon 上达到了接近专家水平的性能。我们尚未了解这种方法的通用性如何，但结果令人鼓舞。

对我们而言，主要的启示是瓶颈并非大语言模型编写 Metal 代码的能力，而是我们提供的上下文和约束条件的质量。我们的 CUDA 翻译层将现有的 NVIDIA 内核专业知识转化为 Apple Silicon 的可操作指导，并让 K-Search 的进化搜索完成其余工作。

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

致谢

这项工作由 IBM 研究团队完成，并基于 UC 伯克利 Sky Lab（Cao 等，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 仓库构建，因此此处的结果可以直接复现。步骤如下：

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

1. 设置凭证

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

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

1. 运行内核搜索

# 在 Apple Silicon 上优化 Flash Attention（世界模型模式）
bash scripts/mac_flash_attention_wm.sh
# 或 Mamba SSM 内核，例如选择性扫描
bash scripts/mamba_selective_scan_fwd_wm.sh

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