K-Search: 將 CUDA 核心專業知識轉移到 Apple Silicon MLX
TL;DR
研究人員已為 K-Search（一個演化核心優化框架）添加結構化翻譯層，該層將數十年的 CUDA 核心專業知識適配到 Apple Silicon 的 MLX 框架。此方法使 AI 能夠自動生成高效能的 GPU 核心，達到與原生 MLX Attention 核心相當的 0.97x 速度，並且相比社群實作，Mamba SSM 核心的預填充（prefill）速度最高可提升 20 倍。
跨平台核心優化的挑戰
編寫高效能的 GPU 核心需要多年專業知識，將這些優化從一家硬體廠商轉移到另一家通常需要從頭重新發現。雖然 CUDA 生態系統擁有針對關鍵操作（如注意力機制和狀態空間模型（SSM））的大量手動調校實作，但較新的生態系統如 Apple Silicon 缺乏此種深度的優化核心，儘管 MLX 框架具備能力，仍常常遺失可觀的性能。
K-Search: 演化核心優化
K-Search 是一個使用迭代迴圈來優化 GPU 核心的演化框架。該流程包含三個主要階段：
- 動作選擇：一個 LLM（在本研究中具體為 Gemini 3.5 Pro Preview）扮演「GPU 核心效能工程師」，分析核心的分類、資料佈局及可能的瓶頸，以從搜尋樹（即「世界模型」）提出優化動作。
- 局部精煉：基於所選動作，程式編寫模型產生候選實作，然後在真實硬體上編譯並進行基準測試。
- 世界模型更新：LLM 根據結果來更新搜尋樹，透過插入新動作、更新優先分數或修剪不成功的路徑。
此搜尋以「Spec」為基礎——一份包含硬體規則與數學約束的領域特定文件，以防止產生無效的原始操作。
CUDA-to-MLX 翻譯層
為了彌合 NVIDIA 與 Apple Silicon 架構之間的差距，研究人員開發了一個翻譯層，該層將 CUDA 概念知識轉換為 MLX/Metal 策略。該層由以下部分組成：
- 概念對照表：一個詞彙表，將 CUDA 原始操作映射到 Metal 對應物，並附帶硬體特定限制（例如，將
__shared__記憶體映射到 Metalthreadgroup記憶體，同時考慮 Apple Silicon 的 32 KB 限制與 NVIDIA 的 48 KB）。 - MLX 特定提示：針對沒有直接 CUDA 對應的模式提供指導，例如使用
simd_shuffle_xor進行以暫存器為基礎的列減少，或使用「exp2 trick」（將 $e^x$ 替換為 $2^{x \log_2 e}$）來利用 Apple 快速的fast::exp2()硬體指令。 - 可重用斷言：將專家核心行為重新框架為演化搜尋必須維持的屬性，而非逐字複製程式碼。
效能基準
注意力核心結果
透過提供完整的翻譯層內容給演化搜尋，研究人員達到了近專家的效能。演化核心獨立發現並實作了以下進階策略：
- Threadgroup 記憶體平鋪
- 線上 softmax
- K-transposition 用於記憶體存取
- exp2 技巧
這使得效能從純演化的 0.26x 提升至 Apple 最先進原生注意力核心的 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 |
|---|---|---|---|
| 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 | 1,092 tok/s | 1,092 tok/s |
| Prefill L=4096 | 6,743 tok/s | 339 tok/s | 1,042 tok/s |
關鍵見解：約 20x 的預填充速度提升歸因於實作了平行（前綴）掃描。而社群的 mlx-lm 實作是依序處理 token，演化核心則使用結合律運算來在 $O(\log N)$ 的依賴步驟中評估序列，充分利用 Apple Silicon GPU 在預填充階段的吞吐量。
未來方向
研究人員正在擴展此工作以支援更多硬體架構，包括 IBM Spyre AIU，並開發更複雜的核心，如融合 MoE 路由和分頁注意力。主要發現是，AI 驅動的核心生成瓶頸不在於 LLM 的編碼能力，而在於提供給模型的架構上下文與約束的品質。