我們正處於運算的新時代。硬體快速演進，不僅GPU速度更快，來自不同供應商的晶片種類也日益增多，每個晶片都有其獨特的架構，並常針對特定AI工作負載進行客製化。軟體變化同樣迅速，幾年前需要數月努力的AI程式碼，現在透過AI編碼工具只需幾分鐘即可生成。

隨著大部分運算都集中在AI上，GPU核心成為其成功的關鍵組成部分。這些是在GPU內部運行的低階程式，要寫出高效能的核心絕非易事，需要多年的專業知識才能做到。將核心從一個供應商的硬體轉移到另一個則更加困難，通常意味著必須從頭重新發現相同的優化方法。

例如，CUDA生態系累積了數十年來之不易的核心專業知識：注意力機制、狀態空間模型（SSM）及其他關鍵操作的手動優化實作，代表著數千小時的工程投入。而較新的硬體生態系（如Apple Silicon、客製化AI加速器等）雖然成長迅速，但缺乏這種深度。

在這項工作中，我們探討了這些專業知識是否能自動轉移。我們以柏克萊Sky Lab的Cao等人提出的K-Search為基礎，這是一個利用AI優化GPU核心的演化式核心搜尋框架，並將其擴展，增加了對MLX的支援。MLX是蘋果為其Apple Silicon晶片設計的機器學習框架。

我們開發了一個新穎的結構化CUDA-to-MLX翻譯層，讓K-Search能將現有的CUDA核心作為知識庫，並將其調整為適用於Apple Silicon的高品質GPU核心，而非從頭開始重建。

我們展示了我們的方法在Apple Silicon上達到了接近專家級的效能，與原生的MLX注意力機制核心相比，速度提升了0.97倍；在Mamba SSM核心上，相較於社群的mlx-lm實作，預填充速度提升了高達20倍。儘管我們主要關注Apple Silicon的MLX核心，但此方法不限於MLX，適用於任何CUDA專業知識可轉移的生態系。

為何選擇MLX？

自2023年末以來，蘋果的MLX框架取得了顯著的採用。隨著Apple Silicon晶片搭載於數億台MacBook和Mac Studio中，MLX使得本地AI推論成為可能，無需雲端成本。其統一記憶體架構對於中型模型（M系列晶片上7B至70B參數）尤其具有吸引力。

然而，在這股發展勢頭之下，存在一個顯著的差距：許多NVIDIA生態系中視為理所當然的效能關鍵核心，例如分頁注意力（paged attention）、優化的SSM掃描核心、融合MoE路由等，在MLX中要麼缺失，要麼未經硬體特定調校而顯得原始。MLX雖然能正確運行模型，但往往犧牲了大量的潛在效能。

這個差距正是本文其餘部分的動機。

K-Search是什麼？

K-Search是一個演化式核心優化框架，最初由我們的第一作者Shiyi Cao在加州大學柏克萊分校Sky Lab開發。給定一個原始核心和硬體規格，它會運行一個迭代優化迴圈：一個大型語言模型（LLM）會推斷接下來要嘗試哪些優化，一個程式碼生成模型會產生候選核心，然後這些候選核心會在真實硬體上編譯並進行基準測試。

測量結果會回饋到搜尋過程中，搜尋會持續精煉，追尋有前景的方向並放棄死胡同，直到效能收斂。搜尋過程由一個「規格」（Spec）作為基礎，這是一個領域特定文件，編碼了硬體規則、優化模式和數學限制，以防止生成的程式碼產生無效原語，並確保候選程式碼能夠實際編譯並高效運行。

在我們的實驗中，單一模型（Gemini 3.5 Pro Preview）同時扮演兩個角色：它維護推理狀態並編寫核心。推理部分被提示為「GPU核心效能工程師」，並被要求在提出任何建議之前，先進行固定分析：分類核心（歸約、掃描、注意力/softmax等）、將參考運算重寫為標準形式、規劃資料佈局和存取模式，並假設在每個運行機制中可能的瓶頸（頻寬、延遲、運算或同步）。只有這樣，它才會發出候選優化方案，每個方案都是一次迭代中可實作的單一變更。

我們將這種持久的推理狀態稱為「世界模型」（world model）。它不是一個扁平的待辦事項列表，而是一棵決策（前綴）樹：每個從根到葉的路徑都構成一個完整的優化計畫，而兄弟分支則是相互競爭的替代方案。每個節點都會被評分——一個介於[0, 10]的總體評分、一個介於[0, 1]的置信度，以及對記憶體頻寬、暫存器壓力、運算/硬體匹配度的節點影響——這樣搜尋就能對部分計畫進行排序並擴展最有前景的計畫。

這棵樹在各輪次中持續存在並成長：精煉一個想法會增加一個子節點，而不是覆蓋其父節點，如果最佳分數在幾輪（停滯期）內未能改善，搜尋就會退回以探索替代分支。

原始的K-Search論文評估了這種搜尋策略在FlashInfer的CUDA核心上的表現。在GQA解碼、MLA解碼、MLA預填充和MoE等任務中，K-Search在相同的120次迭代預算下，比OpenEvolve和ShinkaEvolve更穩定地改進。這些結果確立了我們在此基礎上構建的搜尋框架；本文的其餘部分則探討其優化知識是否能轉移到CUDA之外。

建構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

然而，更有趣的挑戰不僅僅是在MLX上運行K-Search。關鍵的洞察是，專家級的CUDA核心編碼了數十年可轉移到Apple GPU的優化知識，只要你能彌合概念上的鴻溝。僅僅將一個CUDA核心交給LLM並要求它進行移植是不夠的：如果沒有深入的硬體上下文，它會產生語法上有效但架構上錯誤的程式碼（錯誤的分塊大小、無效的原語、不匹配的記憶體假設）。

我們的翻譯層包含：

* 概念對應表：一個結構化的詞彙表，列出CUDA原語及其MLX/Metal對應項，並附帶硬性限制。例如：

* `__shared__` 對應到Metal的執行緒群組記憶體，但有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對應項的操作，提供具體的程式碼級模式，例如在8x8 MMA分塊佈局中使用`simd_shuffle_xor`進行基於暫存器的列歸約，或「exp2技巧」（將$exp(x)$替換為$exp_2(x ext{ log}_2 e)$），以利用蘋果快速的`exp_2`硬體指令加速softmax。

* 可重複使用的斷言：將專家級核心行為重新定義為演化搜尋必須保留的屬性，而非直接複製的程式碼。

匹配專家級核心效能：注意力機制核心

我們評估了Apple Silicon上MLX注意力機制核心的三種配置：（1）一個原始基準線，（2）沒有額外上下文的純粹演化，以及（3）一個完整上下文翻譯層，它為優化器提供了從高效能核心（例如FlashAttention-2）中提取的架構特定實作知識，讓演化搜尋能夠推斷實作策略，而不是從原始核心開始。這三種配置共同讓我們能夠隔離翻譯層的確切影響。

從0.26倍提升到蘋果最先進注意力機制核心速度的0.97倍，說明了翻譯層的重要性。在完整上下文下，演化出的核心獨立發現了FlashAttention 2中的關鍵優化：執行緒群組記憶體分塊、線上Softmax、K轉置以優化記憶體存取，以及exp2技巧。

最後一項將每個Softmax指數替換為以2為底的指數，$e^x = 2^{x ext{ log}_2 e}$，這是精確的，並允許核心直接使用蘋果快速的`fast::exp2()`硬體指令，而無需在運行時支付基數轉換的成本。

預填充速度提升20倍：Mamba SSM核心

為了評估K-Search是否能推廣到注意力機制核心之外，我們將其應用於Mamba使用的狀態空間模型（SSM）核心。與注意力機制不同，其運算瓶頸是遞迴狀態更新而非Softmax，這提供了一個截然不同的優化挑戰。我們在M1 Max上比較了演化實作與社群MLX實作（mlx-lm）和PyTorch參考實作（mamba.py）的效能。

在mamba-370m f16模型上，M1 Max 64GB的評估結果如下：

* 解碼：mlx-mamba（我們）152 token/秒，mlx-lm（社群）116 token/秒。