Reed's News
← 返回精选

K-Search将CUDA内核优化知识迁移到Apple Silicon

AI 75 2026/7/29 2120 字 原文 ↗

图1：CUDA至MLX优化映射表。CUDA的优化经验无需逐条指令照搬，可转化为适配MLX架构的原生优化策略。

我们正身处计算技术的新纪元。硬件迭代日新月异——不仅是GPU性能提升，不同厂商推出的芯片种类日益丰富，每款芯片都有专属架构，且往往针对特定AI工作负载量身打造。软件领域的变革同样迅猛，如今AI编码工具只需几分钟就能完成数年前需耗时数月的开发工作。

当前计算领域的重心已转向AI，GPU内核是AI技术落地的关键组件。这类运行在GPU底层的程序，想要实现高效运行绝非易事，往往需要工程师积累多年经验才能打磨到位。而将内核从某一厂商的硬件迁移到另一平台更是难上加难，通常意味着要从零开始重新摸索优化方案。以CUDA生态为例，它历经数十年积累了大量宝贵的内核优化经验：针对注意力机制、状态空间模型等关键操作的手工调优实现，凝聚了数千小时的工程投入。相比之下，Apple Silicon、定制AI加速器等新兴硬件生态虽发展迅速，却缺乏这样深厚的技术积淀。

本文探讨的核心问题是：能否将CUDA的优化经验自动迁移到其他平台？我们基于加州大学伯克利分校Sky Lab曹世毅等人提出的进化式内核搜索框架K-Search展开研究，为其新增了适配Apple Silicon机器学习框架MLX的后端模块。我们设计了一套创新性的结构化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系列芯片可支持70亿至700亿参数）尤为友好。

然而，在这一发展势头背后，MLX存在显著的性能短板：NVIDIA生态中习以为常的诸多高性能关键内核，如分页注意力、优化后的SSM扫描内核、融合MoE路由等，在MLX中要么缺失，要么仅有无硬件针对性调优的基础版本。MLX虽能正确运行模型，但往往未能充分发挥硬件性能。

正是这一性能短板，成为我们开展后续研究的核心动因。

什么是K-Search？

K-Search是由本文第一作者曹世毅在加州大学伯克利分校Sky Lab开发的进化式内核优化框架。给定一个基础内核和硬件规格，它会启动迭代优化循环：大语言模型（LLM）推理下一步应尝试的优化方向，代码生成模型据此生成候选内核，随后这些候选内核会在真实硬件上完成编译与性能测试。

测试结果会反馈至搜索过程，框架持续优化方案——跟进有潜力的方向，摒弃无效尝试，直至性能趋于稳定。

算法1：基于协同进化世界模型的K-Search。搜索过程交替执行以下步骤：选择最具潜力的优化动作，生成并评估代码直至性能不再提升，通过插入、更新和剪枝操作进化世界模型。改编自Cao等人（2026）。

搜索过程由一份"规格文档（Spec）"约束，这是一个领域专属文档，编码了硬件规则、优化模式和数学约束，可避免生成代码出现无效原语，确保候选内核能够编译并高效运行。

在我们的实验中，单个模型（Gemini 3.5 Pro Preview）承担双重角色：负责推理决策和内核代码编写。推理环节被设定为"GPU内核性能工程师"，在提出优化方案前需完成固定流程的分析：对内核进行分类（归约、扫描、注意力/softmax等），将参考计算重写为标准形式，梳理数据布局与访问模式，推测不同运行场景下的潜在瓶颈（带宽、延迟、计算能力或同步机制）。完成上述分析后，模型才会输出候选优化方案，每个方案对应可在单次迭代中实现的单一修改。

我们将这种持续的推理状态称为世界模型。它并非简单的待尝试事项列表，而是一棵决策（前缀）树：从根节点到叶节点的每条路径代表一套完整的优化方案，同级分支则是相互竞争的备选方案。每个节点都带有评分：范围0-10的overall_rating、范围0-1的confidence，以及对内存带宽、寄存器压力、计算/硬件适配度的impacts评分。借助这些评分，搜索过程可对部分优化方案排序，并优先拓展最具潜力的分支。这棵树会在多轮迭代中持续生长：细化某个想法会添加子节点而非覆盖父节点；若连续几轮最优评分未提升（进入停滞窗口），搜索会回溯并探索其他分支。以下是注意力内核优化过程中生成的一个节点示例：

{
"action": "将线程组内存softmax归约替换为纯寄存器归约：每个SIMD组负责8个查询行，通过simd_shuffle_xor在 lanes间完成归约，移除threadgroup_barrier。",
"difficulty_1_to_5": 4,
"impacts": {
"memory_bandwidth": 8,
"register_pressure": 4, // 风险：若Br>8可能发生寄存器溢出
"compute_hw_fit": 9 // SIMD宽度32；保持8×8分块
},
"overall_rating_0_to_10": 8,
"confidence_0_to_1": 0.7
}

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

图2：K-Search框架概览。框架基于搜索状态$S_t$运行，该状态以搜索树形式组织。树中包含已探索节点（蓝色，如带有对应程序$x_{12}$的已访问状态）和待探索节点边界（橙色，如待验证假设$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论文的核心结果。在三次实验中，针对四个FlashInfer CUDA内核，K-Search在"迄今最佳搜索评分"、各工作负载内核性能及加速比分布上，均优于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。直接让LLM将CUDA内核移植到MLX远远不够：缺乏深层硬件背景的情况下，生成的代码可能语法正确，但架构设计存在缺陷（如分块大小不合理、原语无效、内存假设不匹配）。

我们设计的转换层包含以下核心组件：

  • 概念映射表：结构化的CUDA原语与MLX/Metal等价物对照表，附带硬性约束。例如：

    • __shared__对应Metal的threadgroup内存，但有32KB的硬性容量限制（NVIDIA为48KB）
    • warp_reduce优先对应MMA
    • __syncthreads()替换为threadgroup_barrier(mem_flags::mem_tg)
    • H100的HBM3带宽约3.35TB/s，而M3 Max的统一DRAM带宽约400GB/s，这一差异直接决定了哪些优化策略值得采用。

  • MLX专属提示与模式：针对无直接CUDA等价操作的代码级实现模式，例如在8×8 MMA分块布局中使用simd_shuffle_xor实现基于寄存器的行归约，或利用"exp2技巧"（将$exp(x)$替换为$exp_2(x \log_2 e)$），借助Apple硬件上的快速$exp_2$指令加速softmax计算。

  • 可复用断言：将专家级内核的行为重构为进化搜索必须遵循的属性，而非直接复制代码。

媲美专家级内核性能：Attention内核

我们针对Apple Silicon的MLX注意力内核设置了三种测试配置：(1) 基础基线版本；(2) 无额外上下文的纯进化版本；(3) 完整上下文转换层版本，向优化器提供从高性能内核（如FlashAttention-2）中提取的架构专属实现知识，让进化搜索能够基于已有策略推理，而非从基础内核从头开始。通过对比这三种配置，我们可以精准量化转换层的作用。

图4：Attention内核的性能随优化堆叠的提升曲线。"完整上下文"配置成功探索并实现了双缓冲、循环展开等高级优化策略，达到接近专家级的性能水平。

从Apple最先进注意力内核性能的0.26倍提升至0.97倍，这一飞跃充分体现了转换层的价值。在完整上下文支持下，进化后的内核自主探索出FlashAttention 2中的核心优化手段：线程组内存分块、在线softmax、K转置优化内存访问，以及exp2技巧。其中exp2技巧将所有softmax指数运算替换为以2为底的指数运算，

[e^x = 2^{x \log_2 e},] 该转换完全精确，可让内核直接调用Apple的fast::exp2()硬件指令，无需在运行时进行底数转换，从而节省开销。

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

为验证K-Search能否推广至注意力以外的内核，我们将其应用于Mamba使用的状态空间模型（SSM）内核。与注意力机制不同，SSM的计算瓶颈在于循环状态更新而非softmax，优化挑战截然不同。我们在M1 Max芯片上，将进化后的实现与社区版MLX实现（mlx-lm）及PyTorch参考实现（mamba.py）进行对比。

测试基于mamba-370m（f16精度）、M1 Max 64GB：

指标 mlx-mamba（本文实现） mlx-lm（社区版） mamba.py
解码速度 152 词元/秒 116 词元/秒 40 词元/秒
预填充（序列长度L=512） 5,751 词元/秒 329 词元/秒 1,089 词元/秒
预填充（L=1024） 6,010 词元/秒 327 词元/秒 1,127 词元/秒
预填充（L=2048） 6,612 词元/秒 326 词元/秒 1,092 词元/秒
预填充（L=4096） 6,743 词元/秒 339 词元/秒 1,042 词元/秒

表1：mamba-370m的预填充与解码吞吐量（f16精度，M1 Max 64GB）。本文实现的mlx-mamba预填充吞吐量比社区版mlx-lm高出约20倍，解码性能则相当。

预填充速度提升约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上实现了接近专家级的性能。尽管我们尚未验证该方法的普适性，但现有结果已令人鼓舞。

我们的核心结论是：瓶颈并非LLM编写Metal代码的能力，而是提供给模型的上下文与约束的质量。我们的CUDA转换层将现有的NVIDIA内核优化经验转化为Apple Silicon的可操作指导，剩余工作则由K-Search的进化搜索完成。

目前我们正从多个方向拓展这项研究：支持新架构，当前重点为IBM Spyre AIU及更多硬件开发新内核；新增分页注意力、融合MoE路由等更多内核；优化与K-Search进化循环的集成，让转换上下文的生成更加自动化。

致谢

本研究由IBM Research完成，基于加州大学伯克利分校Sky Lab的K-Search框架（Cao等人，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
  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

完整命令行参考及文档请查看README。