将数十年CUDA内核优化经验自动迁移至Apple Silicon

CUDA生态系统历经数十年积累了大量经过精心调优的GPU内核实现，涵盖注意力机制、状态空间模型等关键操作，背后凝聚了数千小时的工程投入。相比之下，Apple Silicon、定制AI加速器等新兴硬件生态虽然发展迅速，却普遍缺乏这种深度积累。将一个平台的内核移植到另一个平台，往往意味着重新发现相同的优化技巧，效率极低。UC Berkeley与IBM Research的研究团队因此提出：能否让这种专家经验自动完成跨平台迁移？

K-Search是由UC Berkeley Sky Lab的曹世义等人提出的进化式内核优化框架。给定一个初始内核和硬件规格，框架会启动迭代优化循环：大语言模型负责推理下一步应尝试哪些优化，代码生成模型产出候选内核，候选内核随后在真实硬件上编译并基准测试，测量结果反馈回搜索过程，持续精化直至性能收敛。框架的核心是一个被称为世界模型的持久推理状态，以决策前缀树的形式组织，每个节点记录具体优化动作、预估硬件影响、优先级评分和置信度，使搜索能够有效排序并扩展最有潜力的方向。

研究团队为K-Search构建了原生MLX后端，并在此基础上开发了结构化CUDA到MLX翻译层。翻译层的核心包括三部分：概念映射表，将CUDA原语与MLX/Metal等价物对应并标注硬件约束，例如共享内存上限从NVIDIA的48KB降至Metal的32KB，带宽从H100的约3.35 TB/s降至M3 Max的约400 GB/s；MLX专属代码模式，提供无直接CUDA对应物的操作实现方案；以及可复用断言，将专家内核行为转化为进化搜索必须保持的属性，而非直接复制代码。这一设计使模型能够在架构层面理解优化策略，而非进行语法层面的逐行翻译。

在Attention内核评测中，研究团队对比了三种配置：朴素基线、无额外上下文的纯进化搜索，以及包含完整翻译层上下文的配置。朴素基线仅达到原生MLX Attention内核0.26倍的速度，纯进化搜索有所提升但有限，而完整上下文配置最终达到0.97倍，接近专家水平。进化搜索在完整上下文引导下独立发现了FlashAttention 2的核心优化策略，包括线程组内存分块、在线softmax、K矩阵转置内存访问，以及利用Apple硬件快速exp2指令替代标准指数运算的exp2技巧。

在Mamba状态空间模型内核上，K-Search展现出更为显著的性能提升。社区mlx-lm实现未采用并行扫描算法，对状态递推逐token顺序处理，导致Apple Silicon的计算资源大量闲置。研究团队的进化内核实现了并行前缀扫描，将原本O(N)的顺序依赖压缩为O(log N)的并行步骤，在M1 Max上对mamba-370m模型的预填充吞吐量从约329 tok/s提升至约5751至6743 tok/s，实现约20倍加速。解码速度提升相对有限（116 tok/s对比152 tok/s），原因在于单token解码本身无法并行化。

研究团队表示，此次工作的核心发现在于：性能瓶颈并非大语言模型编写Metal代码的能力，而在于提供给模型的上下文质量与约束条件。CUDA翻译层将现有NVIDIA内核专家知识转化为Apple Silicon的可操作指导，再由K-Search的进化搜索完成剩余工作。团队目前正在将该方法扩展至IBM Spyre AIU等更多硬件架构，并计划增加分页注意力、融合MoE路由等更多内核支持，同时持续改进翻译上下文的自动化程度。

要点

• K-Search通过结构化CUDA到MLX翻译层，将NVIDIA生态数十年的内核优化经验自动迁移至Apple Silicon，在Attention内核上达到原生MLX实现0.97倍的速度。
• 在Mamba状态空间模型内核上，进化搜索通过引入并行前缀扫描算法，实现相比社区mlx-lm约20倍的预填充吞吐量提升。
• 翻译层的价值在于架构层面的概念映射，而非逐行代码翻译，使大语言模型能够理解不同硬件的约束差异，如带宽、内存上限等关键参数。
• 该方法不局限于MLX，理论上可推广至任何存在CUDA知识可迁移性的新兴硬件生态，为跨平台内核优化提供了通用框架思路。
• 研究表明，AI驱动的内核优化中，提供给模型的上下文质量与约束条件是决定性能上限的关键因素，而非模型本身的代码生成能力。