在 HotChips 上發表 Jalapeño MLA 核心 以及後續 SemiAnalysis 的評論之後，引發了許多討論。身為 OpenAI 的硬體團隊，我們只是 稍微 觸及了這個小金塊：事實是，AI 正在為我們撰寫核心，而且當它這麼做時，我們並不需要 真的 逐行理解核心在做什麼。我們明顯遺漏了：這樣的事情是如何可能的？與傳統的程式碼生成方法相比，正確的思考方式是什麼？最佳化後的核心與未最佳化的核心一樣可靠嗎？
就我的背景而言，我從事加速器編譯器工作已經超過十年。我創立了 XLA，這是一個優秀的編譯器基礎架構，背後有出色的跨公司團隊和努力。在 OpenAI 的過去兩年多裡，我一直試圖重新構想，在 AI 時代，編譯器應該如何運作。新的編譯器公式將借鑒現有優勢，但不可否認的是，工具箱中有一個強大的新工具可以利用。
這將是一段旅程，但我希望能夠闡明 AI 如何被用於自動化 電腦程式改進，也就是最佳化編譯。我確實認為，透過 AI，我們可能會體驗到我們認為的「編譯器 2.0」。AI 在它能提出的方案上，從根本上受到的限制較少，而它所提出的方案是模型訓練和上下文的結果，這讓我將其歸類為「隨機最佳化器」——這可能會帶來挑戰，但正如我們將看到的，這也是巨大優勢的來源……
大量的學術研究和產業應用已經朝著這個方向前進，並且迅速發掘 AI 參與最佳化編譯器領域的潛力，但我們現在正處於一個需要廣泛解釋的階段。
背景
編譯器接收程式，然後輸出這些程式的翻譯或改進版本。
程式，無論是輸入端還是輸出端，都具有 語意，告訴我們程式的意義、它們可能做什麼，以及如何推理這些它能做的事情。
我們這些從事編譯器工作的人，把它們想像成純函數——它們接收一個資料結構，然後輸出一個應該具有對應語意的資料結構。
有時，我們的編譯器專注於「降低」或「翻譯」。例如，它們可能接收 C 語言，然後輸出 x86-64 組合語言，我們通常認為這是「較低層級」。但實際上，在過程中的一個子部分，它們通常不僅僅是做翻譯……
我們的編譯器，在 實務上，專注於「最佳化」。它們可能接收一個代表程式的資料結構——用我們的行話來說是「中間表示」（IR）——然後嘗試產生該程式的更好版本。有時「更好」意味著執行需要更少的週期，有時意味著它會有更少不必要的程式碼，有時意味著針對我們可以證明「必定為真」的事物進行特化（部分求值）。
現在，簡短地考慮一下，LLM 最初是為了將人類文字從一種語言翻譯成另一種語言而創建的。顯然，翻譯是它們的強項。而且我們可以透過在日常任務中使用 LLM 看到，它們也能夠撰寫新的解決方案並改進現有解決方案。我們許多程式設計師也有要求 LLM「最佳化這段程式碼」的經驗，而且它們確實能做到。（然而，我們需要知道它們正確地最佳化了程式碼，這我們稍後會談到！）這只是為了強調 LLM 具備我們在最佳化編譯器中尋找的能力。
最佳化與最佳性
最佳化編譯器，顧名思義，試圖根據某個目標（通常是執行時間）來 提高 它們正在處理的程式的 最佳性。在一般情況下，對於任意程式來說，這 非常 困難，以至於有一個定理叫做編譯器工程師充分就業定理。（我是在選擇成為編譯器工程師之後才發現這個定理的，但它仍然讓我感到安慰！）
「超級最佳化器」是最佳化編譯器中一個了不起的小子領域。想像有一個給定的程式，我們可以透過語意來說明它的功能。那個 具有相同語意的最最佳化程式是什麼？這就是超級最佳化器試圖解決的問題，而它本質上是一個搜尋問題……
想像我正在嘗試找到具有相同語意的最短程式，並且我有辦法詢問一個候選程式是否具有相同的語意。理論上，我可以按目標順序列舉每個程式，然後選擇具有相同語意的最小程式。
然而，按目標順序列舉每個程式聽起來非常棘手。我最喜歡的一篇學術論文，發表於 2013 年，名為「STOKE」（隨機超級最佳化），提出了這樣的問題：「嗯，如果我們只是隨機地 調整 程式一次又一次，我們最終是否會觀察到最好的程式？」他們提出，透過隨機漫步（以及我們 OG 機器學習朋友馬可夫鏈蒙地卡羅 / Metropolis-Hastings），最終你會看到那個最佳程式。
蒙地卡羅調整通常是愚蠢的（你隨機選擇一個調整），但也很快。LLM 非常聰明（許多推理 token），但相對較慢。
如果，與愚蠢/快速的蒙地卡羅調整不同，我們讓 LLM 找出程式應該前進的方向呢？我們將會有一個 非常 智能的隨機最佳化器，引導我們的程式穿過最佳化的程式空間。
最佳化的直覺
讓我們退一步。想想你認識的那個最能體現「把程式碼最佳化到極致」的人。簡而言之，我們稱他們為「最佳化王 Ollie」。Ollie 可能對哪些類型的程式碼調整可能會有成果有種直覺。Ollie 可能會嘗試一些方法看看是否有效，如果沒有效果，就回滾並嘗試其他方法。但他們對哪些類型的事情是可能的，以及他們 可能 如何擊敗編譯器，有一定的直覺。
Ollie 擁有的這些直覺通常超越了編譯器所能做的。儘管現代最佳化編譯器的結果令人印象深刻，但它們基於相當簡單的規則和啟發式方法。用技術術語來說，它們基於運行到固定點的局部資料流轉換的概念。我們也對考量進行階段排序；也就是說，我們構建編譯器管線來先考慮 A 然後 考慮 B，而不是複合的 AB 問題。排程器和暫存器分配器就是一個臭名昭著的例子，許多博士論文都嘗試過複合排程器-暫存器分配器（以獲得合併階段排序的好處），但在實踐中讓它運作起來一直很有挑戰性。
這就是為什麼 Ollie 的專業知識很有價值。Ollie 通常知道如何用針對特定情況的啟發式方法來平衡幾個 NP 完全問題。因此，它具有更多的特定情境感知和敏感性。Ollie 也能夠應用最佳化編譯器可能無法有效運用的技術，尤其是組合使用時，例如提取、自訂 ABI 或轉換以啟用向量化，或其他讓我們抱怨「我希望編譯器有辦法做到這個……」的事情。
現在考慮一下，AI，透過它擁有的任何推理能力，可能能夠扮演一個迷你 Ollie 的角色。它可能沒有 Ollie 那種關於什麼會成功的匹配直覺，但它對什麼可能是有利的有一點概念，並且它可以進行很多、很多 次的嘗試。
使用這種方法，與 STOKE 論文不同，我們無法 保證 隨著時間趨近無限，我們能看到 最佳 程式，但因為 AI 具有「更像人類」的推理能力，它實際上可以在單位時間內獲得顯著的人類般的進展。
連結回主題：MLA 核心
讓我先說：我 不 知道 AI 為我們的 Jalapeño MLA 核心吐出了什麼低階程式碼，但我 確實 知道如何為 MLA 輸入 numpy 程式碼。
在我之前參與的 XLA 編譯器中，我們會將這些 numpy 操作融合成塊，然後使用一個稱為「發射器」的元程式將其降低為迴圈、指令和低階原始指令。
當 XLA 編譯器 / 發射器程式這樣做時，我不需要關心後面產生了什麼組合語言。對於我們的隨機最佳化器，AI 在概念上取代了發射器元程式——它同時進行降低和最佳化，我們可以要求它進一步最佳化，直到接近理論極限。
我希望這能清楚地說明 AI 的位置，以及它如何類似於現有最佳化編譯器系統中的一個組件。另一個有用的思考方式是：我們認為是「組合語言」的層級現在正在 向上移動。當你輸入普通的 C++ 並以 -O3（典型的最高最佳化級別）編譯它時，你不會期望理解輸出的組合語言，即使你理解了你輸入的 C++。我們在這裡做的是類似的事情，但使用的是更高層級、更數學化的輸入規格。
現在，一個關鍵問題是我們 如何 檢查從 AI 得到的程式是否確實等同於高層級描述 / numpy。這個檢查機制建立了 AI 隨機最佳化過程的可靠性。我預計未來的部落格文章可能會更詳細地討論這個問題，但現在，只要說測試語意等價性是可能的，而且我們正在做，就足夠了。加速器程式特別適合於強大、完整的契約，我們可以驗證「這正是 AI 最佳化後的程式所做的事情」，因為它們在廣泛的背景下是相當數學化和資料流導向的。
請注意，這個領域的許多相關技術是由程式合成子領域的努力所開創的。最佳化編譯器說：「這裡有一個具有語意的程式，讓它變得更好但保持語意等價！」，而程式合成則說：「我相信存在一個具有這些語意的程式，請嘗試找到你能找到的最好的一個」。程式合成是一個比最佳化編譯更困難的問題，但它從根本上受到的約束也更少。這實際上就是像 Ollie 這樣的人，當他們做得比最佳化編譯器更好時所做的事情，而這是 AI 現在可以幫助我們自動化的事情。AI 可以從原始程式中汲取「靈感」，但它不必只對其進行微小的局部轉換。經典的最佳化編譯器 不會 看到「哦，你寫了一個氣泡排序」，並透過理解契約，將其切換為快速排序，但 Ollie 和 AI 都能夠做到這一點。這使得我們更處於隨機最佳化的程式合成機制，而不是經典的最佳化編譯器機制。
所有這些結合在一起，體現在你可以從一個「離 numpy 不遠」的東西開始，等待 48 小時，然後得到一個具有相同語意的最佳化核心，正如我們在 HotChips 演講中所展示的那樣：
正如投影片也提到的，在我們的機器上，我們經常觀察到 AI 的效能超越我們的人類專家，即使 是在我們認為已經調整得相當好的核心上。通常，由於可能需要探索的各種組合 / 排列，仍然存在相當可觀的可達成百分比。對於人類效能工程師來說，這些通常繁瑣得難以處理。
回顧與結論
歸根結底，編譯器只是一個函數。我們將程式交給那個函數，然後得到一個更好的程式版本。我們期望輸出的程式和輸入的程式具有相同的語意。
傳統的最佳化編譯器透過資料流規則和啟發式方法來改進程式。這些方法的來源完全可以理解，但它們能進行的操作也可能更有限。
相比之下，AI，作為一個 隨機 最佳化器，只需要「努力思考」然後吐出結果。它的操作從根本上不受限制，使其更類似於我們的人類專家最佳化器。我們確實需要方法來檢查它吐出的程式是否可靠，並實現我們輸入的相同語意，而我們已經有這些方法。而且這種 AI 最佳化特別適合於具有非常強契約的數學運算。這些契約消除了逐行理解核心在做什麼的需要。
這就是我們如何得到 AI 生成的 MLA 核心！
