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

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

我们正面临计算领域的新纪元。硬件正在快速变化——不仅是更快的 GPU，还有来自不同厂商的、日益多样化的芯片，每种芯片都有其独特的架构，并且通常针对特定的 AI 工作负载进行了定制。软件的变化同样迅速，AI 编码工具现在能在几分钟内生成几年前需要数月努力才能完成的工作。由于如今计算领域如此集中于 AI，GPU kernel 是其成功的关键组成部分。这些是在 GPU 内部运行的低级程序，编写高效的 kernel 绝非易事——需要多年的专业知识才能做好。将一个 kernel 从一家厂商的硬件迁移到另一家厂商的硬件更加困难，并且通常意味着从头重新发现相同的优化方法。例如，CUDA 生态系统已经积累了数十年来之不易的 kernel 专业知识：手动调优的 attention、状态空间模型以及其他关键操作的实现，代表了数千小时的工程投入。较新的硬件生态系统（Apple Silicon、定制 AI 加速器等）发展迅速，但缺乏这种深度。在这项工作中，我们探究这种专业知识能否被自动迁移。

我们基于 K-Search 进行构建，这是由加州大学伯克利分校 Sky Lab 的 Cao 等人提出的一个进化式 kernel 搜索框架，它使用 AI 来优化 GPU kernel。我们为其扩展了一个面向 MLX 的后端——MLX 是 Apple 为其 Apple Silicon 芯片开发的机器学习框架。我们开发了一种新颖的结构化 CUDA 到 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？ Apple 的 MLX 框架自 2023 年底以来得到了显著采用。随着 Apple Silicon 被应用于数亿台 MacBook 和 Mac Studio 中，MLX 使得无需云成本的本地 AI 推理成为可能。其统一内存架构使其对中等规模模型（M 系列芯片上 7B–70B 参数）特别有吸引力。然而，在这一发展势头之下，存在一个显著的差距：许多 NVIDIA 生态系统视为理所当然的性能关键型 kernel——如 paged attention、优化的 SSM scan kernel、融合的 MoE routing——要么缺失，要么在没有硬件特定调优的情况下实现得过于简单。MLX 能够正确运行模型，但通常留下了大量性能潜力未被挖掘。这一差距正是推动本文后续内容的原因。

什么是 K-Search？ K-Search 是一个进化式 kernel 优化框架，最初由我们的第一作者 Shiyi Cao 在加州大学伯克利分校 Sky Lab 开发。给定一个朴素 kernel 和一份硬件规格说明，它会运行一个迭代优化循环：一个 LLM 推理下一步应该尝试哪些优化，一个代码编写模型生成候选 kernel，然后这些候选 kernel 在真实硬件上被编译和基准测试。测量结果反馈回搜索过程，搜索过程不断优化，追求有希望的方向并放弃死胡同，直到性能收敛。

算法 1：通过共同进化的世界模型进行 K-Search。 搜索过程在选择最有希望的行动、实例化和评估代码直到改进停滞，以及通过插入、更新和剪枝操作来进化世界模型之间交替进行。改编自 Cao 等人 (2026)。搜索过程由一个 Spec（规格说明）指导：一份领域特定的文档，编码了硬件规则、优化模式和数学约束，这防止了生成的代码产生无效的原语幻觉，并确保候选代码能够实际编译并高效运行。

在我们的运行中，一个单一模型（Gemini 3.5 Pro Preview）扮演了两个角色：它维护推理状态并编写 kernel。推理部分被提示为“GPU kernel 性能工程师”，并被要求在提出任何建议之前完成一个固定的分析流程：对 kernel 进行分类（reduction、scan、attention/softmax……），将参考计算重写为规范形式，规划出数据布局和访问模式，并假设在每个运行时阶段可能的瓶颈（带宽、延迟、计算或同步）。只有在这之后，它才会生成候选优化方案，每个方案都是一个可以在一次迭代中实现的单一更改。我们将持久的推理状态称为世界模型。它不是一个扁平的待尝试列表，而是一个决策（前缀）树：每个根到叶子的路径组成一个完整的优化计划，兄弟分支是相互竞争的替代方案。每个节点都被评分——一个 [0, 10] 范围内的 overall_rating，一个 [0, 1] 范围内的 confidence，以及每个节点对内存带宽、寄存器压力和计算/硬件适配度的影响——这样搜索就可以对部分计划进行排序，并扩展最有希望的计划。这棵树在多次迭代中持续存在并增长：完善一个想法会添加一个子节点，而不是覆盖其父节点，如果最佳分数在几轮迭代（停滞窗口）内未能改善，搜索会回退以探索替代分支。一个在 attention kernel 搜索过程中出现的节点示例如下：

{
  "action" : "将线程组内存 softmax reduction 替换为仅寄存器 reduction：每个 SIMD 组拥有 8 个查询行，并使用 simd_shuffle_xor 在通道间进行 reduction，移除一个 threadgroup_barrier。",
  "difficulty_1_to_5" : 4,
  "impacts" : {
    "memory_bandwidth" : 8,
    "register_pressure" : 4, // 风险：如果 Br > 8 则发生溢出
    "compute_hw_fit" : 9 // SIMD 宽度 32；保持 tile 8 x 8
  },
  "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}$）组成的 Frontier（前沿）。工作流程迭代三个阶段：(1) 行动选择，根据世界模型估计的优先级分数 $V$ 从前沿中检索最有希望的行动节点；(2) 局部细化，随机策略 $\pi_{\mathrm{code}}$ 采样具体实现直到停滞；以及 (3) 世界模型更新，LLM 对轨迹进行推理，通过插入（添加新行动）、更新（调整 $V$，例如 $u_{11}$ 从 0.9 降至 0.6）和剪枝（移除不太有希望的节点如 $u_{10}$）来更新搜索树。

原始的 K-Search 论文在来自 FlashInfer 的 CUDA kernel 上评估了这种搜索策略。在 GQA decode、MLA decode、MLA prefill 和 MoE 上，K-Search 在相同的 120 次迭代预算下比 OpenEvolve 和 ShinkaEvolve 取得了更一致的改进。这些结果确立了我们在本文中构建的搜索框架；本文的其余部分探讨其优化知识能否迁移到 CUDA 之外。

图 3：原始 K-Search 论文的主要结果。 在三次运行中，K-Search 在四个 FlashInfer CUDA kernel 上，在迄今为止最佳搜索分数、每个工作负载的 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 编译和执行。
• 更新后的 kernel 生成器提示，用于编写和修改 Metal/MLX kernel。
• 使用 mlx.core 测量工具的 MLX 特定基准测试集成。

将 CUDA 专业知识翻译到 MLX

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

我们的翻译层包括：

• 概念映射表：一份结构化的词汇表，将 CUDA 原语与其 MLX/Metal 对应项进行映射，并附有硬性约束。例如：
  • __shared__ 映射到 Metal 的 threadgroup memory，但有 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 进行基于寄存器的行 reduction，或者使用“exp2 技巧”（将 $exp(x)$ 替换为 $exp_2(x \log_2 e)$）以在 Apple 快速的 exp2 硬件指令上实现更快的 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 memory tiling、online 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，这提供了一个截然不同的优化挑战。我们将进化后的实现与社区 MLX 实现（mlx-lm）以及 PyTorch 参考实现（mamba.py）在 M1 Max 上进行了比较。

在 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（我们的）的 prefill 吞吐量比社区 mlx-lm 基线高出约 20 倍，而 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)$ 个依赖步骤中求值，而不是 $O(N)$。mlx-lm 跳过了这一点，一次只处理一个 token，使得 Apple Silicon 的大部分计算能力闲置；我们进化后的 Metal kernel 应用了扫描，从而更充分地利用了 GPU 的吞吐量。这种增益体现在 prefill 上，因为整个序列可用于并行扫描，而不体现在单 token decode 上，因为每一步只有一个新 token，没有扫描可以并行化——这就是为什么 decode 行大致持平而 prefill 快了约 20 倍。mamba.py 在 prefill 和 decode 上都慢，因为它是一个 PyTorch 参考实现，在 Apple Silicon 上会回退到 CPU 或 MPS，放弃了 MLX 的 Metal 后端所能实现的硬件特定优化。

下一步是什么？

在我们研究的两个 kernel 上，基于结构化跨平台翻译知识的 AI 驱动进化式 kernel 搜索，在没有 GPU 专家团队从头开始的情况下，在 Apple Silicon 上达到了接近专家水平的性能。我们尚不清楚这能泛化到何种程度，但这一结果令人鼓舞。对我们来说，主要的收获是瓶颈不在于 LLM 编写 Metal 代码的能力，而在于我们提供给它的上下文和约束的质量。我们的 CUDA 翻译层将现有的 NVIDIA kernel 专业知识转化为适用于 Apple Silicon 的可操作指导，并让 K-Search 的进化搜索完成其余工作。

我们正在几个方向上积极扩展这项工作：支持新架构，当前工作重点是为 IBM Spyre AIU 和更广泛的硬件目标开发新 kernel；添加更多 kernel，例如 paged attention 和融合的 MoE routing；以及改进与 K-Search 进化循环的集成，使翻译上下文更加自动化。

致谢

这项工作由 IBM Research 进行，并基于加州大学伯克利分校 Sky Lab 的 K-Search（Cao 等人，2026）。我们欢迎来自 MLX 和更广泛的 AI 系统社区的协作与反馈。如果您正在从事非 CUDA 硬件的 kernel 优化工作，我们很乐意听取您的意见。

引用

@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. 运行 kernel 搜索

   # 在 Apple Silicon 上优化 Flash Attention（世界模型模式）
   bash scripts/mac_flash_attention_wm.sh
   
   # 或者 Mamba SSM kernel，例如 selective scan
   bash scripts/mamba_selective_scan_fwd_wm.sh
   

完整的 CLI 参考和文档在 README 中。