K-Search: 将 CUDA 内核专业知识迁移到 Apple Silicon MLX

TL;DR

研究人员在 K-Search（一个进化内核优化框架）上添加了一个结构化的翻译层，该层将数十年的 CUDA 内核专业知识适配到 Apple Silicon 的 MLX 框架。此方法使 AI 能够自动生成高性能 GPU 内核，达到原生 MLX Attention 内核速度的 0.97x，并且相比社区实现，Mamba SSM 内核的预填充速度提升最高可达 20x。

跨平台内核优化的挑战

编写高效的 GPU 内核需要多年的专业知识，并且将这些优化从一个硬件供应商迁移到另一个通常需要从头重新发现它们。虽然 CUDA 生态系统拥有针对关键操作（如注意力机制和状态空间模型（SSM））的大量手工调优实现，但较新的生态系统如 Apple Silicon 缺乏这种深度的优化内核，尽管 MLX 框架具备能力，但往往会留下大量性能未被利用。

K-Search: 进化内核优化

K-Search 是一个使用迭代循环来优化 GPU 内核的进化框架。该过程包括三个主要阶段：

1. 行动选择：一个 LLM（在本研究中具体为 Gemini 3.5 Pro Preview）充当“GPU 内核性能工程师”，分析内核的分类、数据布局和可能的瓶颈，以提出来自搜索树（即“世界模型”）的优化动作。
2. 局部细化：基于所选动作，代码编写模型生成候选实现，然后在真实硬件上进行编译和基准测试。
3. 世界模型更新：LLM 基于结果推理，通过插入新动作、更新优先级分数或修剪不成功的路径来更新搜索树。

此搜索以“Spec”为基础——一份包含硬件规则和数学约束的领域特定文档，以防止生成无效的原语。

CUDA-to-MLX 翻译层

为了弥合 NVIDIA 和 Apple Silicon 架构之间的差距，研究人员开发了一个翻译层，该层将 CUDA 概念知识转换为 MLX/Metal 策略。该层由以下部分组成：

• 概念映射表：一个词汇表，将 CUDA 原语映射到 Metal 等价物，并伴随硬件特定的约束（例如，将 __shared__ 内存映射到 Metal threadgroup 内存，同时考虑到 Apple Silicon 上的 32 KB 限制与 NVIDIA 上的 48 KB 限制）。
• MLX 特定提示：针对没有直接 CUDA 等价物的模式提供指导，例如使用 simd_shuffle_xor 进行基于寄存器的行规约，或使用“exp2 技巧”（将 $e^x$ 替换为 $2^{x \log_2 e}$）以利用 Apple 快速的 fast::exp2() 硬件指令。
• 可重用断言：将专家内核行为重新框架为进化搜索必须维持的属性，而不是简单地复制代码。

性能基准

注意力内核结果

通过为进化搜索提供翻译层的完整上下文，研究人员实现了接近专家级的性能。进化内核独立发现并实现了以下高级策略：

• Threadgroup 内存分块
• 在线 softmax
• 用于内存访问的 K-转置
• exp2 技巧

这使得性能从纯进化的 0.26x 提升到苹果最先进的原生注意力内核速度的 0.97x。

Mamba SSM 内核结果

K-Search 被应用于 Mamba 状态空间模型（SSM）内核以测试泛化能力。在使用 mamba-370m f16 的 M1 Max (64GB) 上，进化的 mlx-mamba 内核相比社区的 mlx-lm 实现展示了预填充吞吐量的巨大提升：

指标	mlx-mamba (ours)	mlx-lm (community)	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	1,092 tok/s	1,092 tok/s
预填充 L=4096	6,743 tok/s	339 tok/s	1,042 tok/s

关键洞察：~20x 预填充加速归因于实现了并行（前缀）扫描。虽然社区的 mlx-lm 实现按顺序处理标记，但进化内核使用关联结合在 $O(\log N)$ 依赖步骤中评估序列，充分利用了 Apple Silicon 在预填充阶段的 GPU 吞吐量。

未来方向

研究人员正在扩展此工作以支持更多硬件架构，包括 IBM Spyre AIU，并开发更复杂的内核，如融合 MoE 路由和分页注意力。主要发现是，AI 驱动的内核生成的瓶颈不在于 LLM 的编码能力，而在于提供给模型的架构上下文和约束的质量。