宅中地 - 每日更新
宅中地 - 每日更新

贊助商廣告

X

斯坦福與亞馬遜聯手打造AI晶片"自動調音師":讓大模型自己學會優化底層代碼

2026年04月27日 首頁 » 熱門科技

這項由斯坦福大學電腦科學系與亞馬遜雲服務(AWS)聯合開展的研究,發表於2026年第九屆MLSys機器學習與系統國際會議(MLSys 2026),會議地點位於美國華盛頓州貝爾維尤。論文編號為arXiv:2511.15915v2,發布於2026年4月15日,感興趣的讀者可通過該編號檢索完整論文。

一、為什麼一塊晶片發揮不出應有的力量

每當我們使用ChatGPT、Gemini或任何大型AI工具時,背後都有成千上萬塊專用晶片在瘋狂運轉。這些晶片被稱為AI加速器,它們就像一台台精密的樂器,天生具備演奏高難度樂章的能力。然而,樂器再好,如果演奏者不懂得如何駕馭它,也只能彈出平庸的音符。

在AI晶片的世界裡,這個"演奏者"就是所謂的**核心程序**(kernel)——一段專門告訴晶片如何處理矩陣乘法、注意力機制等計算任務的底層代碼。核心寫得好,晶片可以發揮出接近理論極限的性能;核心寫得差,再貴的晶片也只能跑在30%甚至更低的效率上。問題在於,針對每一塊新晶片寫出高質量的核心,需要具備深厚專業知識的硬體工程師花費大量時間和精力去摸索調試。英偉達斯坦福與亞馬遜聯手打造AI晶片自動調音師讓大模型自己學會優化底層代碼H100晶片於2022年發布,工程師們花了將近一年時間才把注意力機制的核心調到理論峰值的37%,又花了一年多才勉強接近85%。

亞馬遜自研的Trainium晶片是這個故事的主角背景。它是一款專門用於訓練大型AI模型的加速器,擁有獨特的硬體架構,對應的編程接口叫做神經元核心接口(NKI,Neuron Kernel Interface)。由於這套接口相對較新,工程師們缺乏成熟的調優經驗和現成的優化套路,很多核心程序都還運行在遠低於硬體極限的水平上。

研究團隊注意到這個困境,提出了一個大膽的想法:既然大型語言模型(LLM,就是ChatGPT那類會寫文章、能聊天的AI)在理解和生成代碼方面已經頗有造詣,能不能讓AI來幫AI晶片調音?更進一步,能不能讓這個AI調音師在工作中不斷積累經驗,越調越好,最終不需要任何人工干預就能自主掌握優化技巧?

這就是AccelOpt(Accelerator Optimizer,加速器優化器)誕生的初衷。

二、AccelOpt是什麼:一位會自我進化的調音師

以調鋼琴為比喻來理解AccelOpt會非常直觀。一位頂尖鋼琴調音師第一次面對一台從未見過的新型鋼琴,他會怎麼做?他會先試彈幾個音,聽聽哪裡跑調了,再根據經驗判斷可能的原因,然後動手調整,反覆驗證。更重要的是,每次成功調好一台琴,他都會把這次的經驗記錄下來——"這個型號的琴,中音區的弦往往偏緊,要用某種特定的手法處理"——下次遇到類似的琴,他就能更快更准地找到問題所在。

AccelOpt做的事情與此如出一轍。它面對的"鋼琴"是Trainium晶片上運行的NKI核心程序,面對的"調音任務"是讓這些程序跑得更快、更接近晶片的理論極限。整個系統由三個相互協作的AI智能體組成,分別扮演不同的角色,再加上一套"束搜索"機制和一個持續更新的"優化記憶庫",共同構成了這台自動調音機器。

三個智能體各司其職,形成一條完整的調音流水線。**規劃師**(Planner)負責分析當前核心的運行剖析數據——也就是晶片在執行這段代碼時,哪個部分耗時最長、哪裡在白白浪費資源——然後提出一個具體的優化方向。**執行師**(Executor)拿到規劃師的方案後,動手將其轉化為實際的代碼改動,真正生成一個新版本的核心程序。**總結師**(Summarizer)則在一次成功的優化完成之後,把這次的經驗提煉成可以復用的通用規律,存入記憶庫,供未來的任務參考。

三、束搜索:不把寶押在一次嘗試上

光有三個智能體還不夠。一個經驗豐富的調音師知道,面對一台複雜的鋼琴,光靠一次嘗試找到最優解幾乎是不可能的。他可能需要嘗試多種不同的調整思路,分頭推進,看看哪條路走得最順。

AccelOpt採用的**束搜索**(beam search)機制正是這個道理。在每一輪疊代中,系統不是只保留一個最新的核心版本,而是同時維護著一批表現最好的候選核心——就像調音師同時在腦子裡保留著幾套並行的方案。具體來說,系統會給每個候選核心生成若干條不同的優化方向(由規劃師提出),再對每個方向嘗試若干次代碼實現(由執行師完成),最後從所有這些新生成的版本中,挑選出表現最佳的一批,作為下一輪疊代的起點。

這套機制的好處在於,它避免了"走進死胡同"的風險。如果某次優化嘗試走錯了方向,損失的只是其中一個分支,其餘分支還在繼續前進。隨著疊代輪次的增加,整個候選集的整體質量會穩步攀升。研究團隊的實驗也驗證了這一點:束搜索比單純地重複採樣同一個提示詞效果要好得多,因為它真正實現了從優秀走向更優秀,而不是在原地反覆打轉。

四、優化記憶庫:經驗比黃金更值錢

束搜索解決了"怎麼探索"的問題,但還有另一個挑戰:AI系統在這次優化任務中學到的東西,能不能在下一次任務中發揮作用?

這就是**優化記憶庫**(optimization memory)的用武之地。它本質上是一個動態更新的經驗檔案,記錄著每次成功或失敗的優化嘗試中,哪段代碼的哪種改動帶來了顯著的速度提升或下降。更重要的是,總結師會把這些具體的代碼改動抽象成更通用的規律,比如"把不變的矩陣轉置操作移到循環外面,可以避免重複計算"。

記憶庫的設計有幾個值得關注的細節。一方面,它不僅記錄"成功案例"(慢變快),也記錄"失敗案例"(快變慢),因為知道什麼不該做有時候和知道什麼該做同樣重要。另一方面,為了防止記憶庫被重複的、相似的經驗刷屏,系統在選取經驗條目時會刻意保持多樣性——不同的候選核心、不同的優化方向,各自貢獻一個最有代表性的經驗點。

記憶庫有兩個關鍵參數控制其行為:ExpN控制庫的容量(能保存多少條歷史經驗),TopK控制每輪疊代能往庫里新增多少條新經驗。研究團隊發現,增大容量(ExpN)比增大每輪更新量(TopK)更划算——保留更多歷史沉澱的價值,高於在每輪疊代中急著塞入更多新經驗。這就好像一位調音師,積累十年工齡的深厚閱歷,比剛入行時每天快速記筆記要更有價值。

五、NKIBench:給調音師設計一套考題

研究團隊在評估AccelOpt時遇到了另一個問題:沒有合適的標準考題。已有的AI晶片核心基準測試,要麼不包含NKI程序,要麼只是簡單地比較"優化後比優化前快了多少倍",卻無法告訴你這個核心距離晶片的極限還有多遠的路要走。

於是研究團隊自己動手構建了**NKIBench**,這是第一個專門針對Trainium晶片NKI核心的基準測試套件。所有14個測試任務都來自真實的大模型工作負載,包括DeepSeek斯坦福與亞馬遜聯手打造AI晶片自動調音師讓大模型自己學會優化底層代碼、Qwen3、Falcon等主流模型中涉及的矩陣乘法、批量矩陣乘、注意力機制、RMSNorm歸一化、LoRA微調、Mamba序列模型等關鍵算子。

NKIBench的一個重要創新在於,它為每個任務計算了硬體理論峰值性能,並用實際運行時間占峰值性能的百分比來衡量核心的優化程度。這就好比一輛賽車在跑圈,不是只看它比上一版車快了幾秒,而是直接告訴你它已經跑到了這條賽道理論最快圈速的多少百分比。這個指標更能客觀反映優化工作還有多大的提升空間。

理論峰值的計算基於晶片的三個核心瓶頸:HBM頻寬(晶片從外部內存讀寫數據的速度)、張量引擎算力(專門做矩陣乘法的計算單元的速度)、向量引擎算力(做其他數學運算的計算單元的速度)。具體公式是取這三個瓶頸中最緊的那個作為理論最快速度,然後用實際延遲與之相比。

六、調音的成果:從49%到61%

經過系統性的測試,AccelOpt的表現相當令人滿意。在Trainium 1晶片上,NKIBench測試集的平均峰值吞吐量占比從優化前的49%提升到了61%;在更新的Trainium 2晶片上,則從45%提升到59%。換句話說,這批核心程序在優化後,能榨取出晶片更多的潛力,平均提升幅度在12到14個百分點之間。

更有趣的是與頂尖商業模型的對比。Anthropic公司的Claude Sonnet 4是當前公認的最強代碼生成模型之一,研究團隊也用它來做同樣的核心優化任務(採用重複採樣的方式,不斷生成新版本直到預算用完)。最終,AccelOpt在使用開源模型(Qwen3-Coder-480B作為執行師,gpt-oss-120b負責其餘角色)的情況下,達到了與Claude Sonnet 4相近的優化效果,但花費的API調用費用僅為後者的二十六分之一。

這個對比背後有一個很有意思的發現:Claude Sonnet 4的"重複採樣"策略,本質上是在同一個起點反覆嘗試,就像一個調音師一遍遍地試彈同一個音,期望某次偶然能找到最優解。而AccelOpt的束搜索則是從優秀的版本出發繼續精進,每一步都站在前一步的肩膀上。這種疊代式進化的策略,讓開源模型在成本上大幅領先的同時,仍能保持競爭力。

七、AccelOpt發現了哪些具體的優化技巧

研究團隊特別列舉了AccelOpt自主發現的幾類典型優化,展示了這套系統的實際能力深度。

第一類是**窺孔優化**(peephole optimization),也就是對局部代碼片段進行精簡和替換。比如,AccelOpt會自動識別出`θ_t–1 – γλθ_t–1`可以化簡為`(1–γλ)θ_t–1`,從而把一次減法和一次乘法合併成一次乘法,減少計算量。它還能識別`reciprocal(sqrt(...))`這個組合可以被替換為更高效的`rsqrt(...)`單一指令,減少中間臨時變量的創建。對於SiLU激活函數,AccelOpt發現`x/(1+e^{-x})`可以改寫為`x·sigmoid(x)`,從而調用NKI專門優化過的sigmoid指令,獲得更高效率。

第二類是**循環變換優化**,這類優化需要更深層次的推理。研究團隊展示了一個BatchMatmul加Softmax融合算子的優化案例,非常能說明問題。初始版本的核心因為中間變量`v`和`p`需要跨越兩個循環存活,被迫溢出到晶片外部的慢速內存(HBM)中,造成大量額外的數據搬運開銷,延遲高達12毫秒,向量引擎利用率只有46%。AccelOpt在第一步發現了這個溢出問題,提出重新計算`v'`來消除溢出,延遲降到8.2毫秒,但代價是引入了額外的矩陣乘法計算。系統沒有滿足於這個結果,繼續推進,在下一步徹底重構了循環結構,去掉了重新計算和多餘的外層循環,最終版本延遲降到6.4毫秒,向量引擎利用率躍升至84%。

第三類是**循環不變量外提**(loop invariant code motion)。這個概念用大家熟悉的場景來說:假設你要做20道數學題,每道題開始都要先查一遍同一張乘法表——合理的做法顯然是只查一次,把結果記在手邊。AccelOpt在Transpose+Matmul的核心中發現,LHS矩陣的轉置操作在每次外層循環疊代中都被重複執行了16次,而這個轉置結果其實從不改變。系統自動將這個操作提到循環外面,緩存在全局緩衝區里,消除了冗餘的15次重複計算。

八、哪裡調不動:AccelOpt的局限性

公平的評估不僅要看成績,也要誠實面對局限。研究團隊通過細緻的觀察,總結出AccelOpt遭遇"停滯"的兩類典型情況。

第一類是"已經夠好了"的停滯。當一個核心在經過若干輪優化後,已經達到了硬體峰值的80%以上,剩下的提升空間自然越來越小。在Matmul(K=5120)這個測試案例中,系統在第7次疊代後達到約82%峰值吞吐量後便不再有顯著進展——但這並非因為系統停止了探索,從後續疊代的數據變化可以看出,系統仍在嘗試各種不同策略,只是已經沒有足夠的提升餘地了。這類停滯是"勝利後的停滯",本質上是好事。

第二類是"無從下手"的停滯。BatchMatmul(B=16, K=64)這個案例是最典型的例子。這個算子的問題規模很小,所有數據都能放進晶片的片上緩存中,因此內存傳輸效率天然就接近100%,沒什麼可優化的餘地。更棘手的是,它的縮減維度K=64隻有晶片張量引擎原生支持的維度128的一半,導致硬體利用率很難提上去,當前的NKI接口也缺乏應對這種情況的有效工具。在這個案例的整個優化過程中,幾乎所有的性能指標都紋絲不動,系統生成的大多數改動要么正確率極低,要麼性能不見起色。這類問題提示了NKI接口本身的設計局限,也指出了未來可以改進的方向。

九、與專家對決:AI能超過人工調優嗎

研究團隊還做了一個很有說服力的比較實驗,把AccelOpt的成果與人類專家手工優化的版本放在一起評判。

對於Mamba序列模型的核心,NKI官方教學提供了三個由人工逐步優化的版本,分別達到28.4%、30.1%和52.7%的峰值吞吐量占比。AccelOpt從同一個起點(28.4%的版本)出發,自主優化到了54.6%,略高於人類專家的最佳版本(52.7%)。更有趣的是,AccelOpt生成的最優核心採用了與人類專家完全不同的循環順序,走的是另一條路,卻到達了更高的終點。

對於RoPE位置編碼的核心,NKI官方示例提供了一個達到21.1%峰值的參考版本。AccelOpt從這個版本出發,將其提升到29.6%,實現了1.4倍的加速。

研究團隊指出,AccelOpt超越人類專家的根本原因在於**規模效應**:人類工程師通常只能串行地逐個優化核心,而AccelOpt可以同時並行探索大量不同的優化方向。當人類工程師還在仔細思考某個方向是否值得深入時,AccelOpt已經同時跑完了幾十個並行實驗。

此外,研究團隊還把AccelOpt應用到了斯坦福大學CS149並行計算課程的教學中。在2025年秋季學期,他們用AccelOpt優化了一個課程外的二維卷積(Conv2D)核心,從去年的參考實現(9.54%峰值吞吐量)提升到了48.8%。基於AccelOpt提出的優化思路,課程組設計了一道附加題,131支學生隊伍中有33.6%成功完成了挑戰。學生們在這個過程中具體學到了兩個原則:把時序上順序執行的計算轉化為空間上並行執行,以及針對特定工作負載在硬體約束下進行專門化設計。

十、成本控制:聰明地花每一分錢

AccelOpt是一個需要持續調用AI模型接口的系統,因此成本控制是不得不考慮的現實問題。研究團隊對各種參數組合和模型選擇做了細緻的成本收益分析。

在模型選擇方面,一個出人意料的發現是:換用不同的規劃師模型(無論是小巧的gpt-oss-20b還是強大的Qwen3-235B-Thinking),對最終優化效果幾乎沒有顯著影響;但換用不同的執行師模型,效果差距則非常明顯。這說明,在AccelOpt的工作流程中,正確理解並實現代碼改動的能力(執行師的任務)遠比制定優化方向的能力(規劃師的任務)更關鍵。因此,提升系統性能應該優先考慮用更強的執行師,而不是花錢升級規劃師。

在內存配置方面,研究團隊發現增大ExpN(擴大記憶庫容量、保留更多歷史經驗)比增大TopK(每輪疊代更激進地更新記憶)更具性價比。以Qwen3-Coder-30B為執行師為例,將ExpN從8增加到16,額外花費12.33美元換來了4.6%的加速提升;而gpt-oss-120b在同樣條件下,額外花費13.81美元只帶來了0.6%的提升,顯示不同模型對擴充歷史經驗的敏感度差異很大。

在整體對比中,使用Claude Sonnet 4做重複採樣的總花費約為5806美元,最終平均加速比約1.222倍;而AccelOpt使用gpt-oss-120b,花費約139美元,達到1.235倍加速。成本差距高達四十倍,效果還略有反超。

十一、不只是Trainium:通用性展望

AccelOpt的設計從一開始就考慮了通用性。束搜索機制和優化記憶庫的核心設計與具體硬體無關,要把系統遷移到新的平台,只需要兩樣東西:一是針對新平台的性能剖析服務,二是包含新平台基礎知識的提示詞模板。

研究團隊已經在GPU平台上做了初步驗證。他們將AccelOpt應用於FlashInfer-Bench中的24個Triton核心(運行在英偉達H100 GPU上),使用gpt-oss-120b實現了平均1.27倍的加速,其中對GQA解碼核心的峰值加速達到3.19倍。研究團隊認為,由於大型語言模型在訓練數據中見過大量GPU編程的示例,AccelOpt在GPU上的效果應該比在相對陌生的Trainium上更好,這一初步結果也印證了這個判斷。

說到底,AccelOpt做的事情是把"讓AI晶片發揮最大效能"這件原本需要頂尖專家耗費數年時光的工作,部分自動化了。它不是萬能的——面對某些特殊情況它也會束手無策,也還無法跨越硬體接口本身的設計限制——但它確實證明了一件事:AI可以通過自我探索和經驗積累,在一個全新的、陌生的技術領域裡,逐步逼近甚至偶爾超越人類專家的水平,而且代價可以比人工低得多。

對於普通用戶來說,這意味著未來大模型的推理速度可能更快、運行成本更低,因為驅動這些模型的底層代碼會越來越高效。對於AI工程師來說,這意味著他們不再需要在每款新晶片上從頭開始漫長的調優之旅,而是可以把這項繁重的初步工作交給系統自動完成,然後在系統探索的基礎上做更有創意的工作。

AccelOpt和NKIBench的代碼與數據已經開源,這14個測試任務也作為了一個起點,供未來的研究者繼續深入探索。對這個方向感興趣的讀者,可以通過arXiv編號2511.15915v2查閱完整論文,以及在GitHub上訪問zhang677/AccelOpt項目獲取代碼。

---

Q&A

Q1:AccelOpt的優化記憶庫和普通AI的上下文窗口有什麼區別?

A:普通AI的上下文窗口是單次對話的臨時記憶,對話結束後清空。AccelOpt的優化記憶庫則是跨輪次持續積累的經驗檔案,專門儲存那些經過實際硬體驗證、帶來真實速度提升的代碼改動規律。它會被整理成通用的優化原則和代碼片段對,並在之後的每一輪優化中作為參考提供給規劃師智能體,相當於一本通過實戰不斷擴充的經驗手冊,而不是每次從零開始。

Q2:AccelOpt能用來優化普通Python或C++代碼嗎?

A:目前AccelOpt的設計專門針對AI加速器的核心程序(如Trainium的NKI程序和GPU的Triton程序),它的優化邏輯依賴硬體性能剖析數據(如內存頻寬占用、引擎利用率等),以及對特定硬體架構的專業知識。不過系統的核心框架——束搜索加上自我更新的經驗記憶——理論上可以遷移到其他代碼優化場景,只需要替換對應的性能評估工具和領域知識提示詞。

Q3:NKIBench中哪類核心最難優化?

A:從測試結果來看,規模較小、數據能完全裝入片上緩存的核心最難通過代碼層面的優化來提升性能,比如BatchMatmul(B=16, K=64)這個案例。這類問題的瓶頸不在於數據搬運或計算冗餘,而在於硬體規格本身的限制——例如Trainium的張量引擎原生支持128維度的縮減,而K=64隻有這個維度的一半,導致硬體利用率天花板極低。AccelOpt在這類情況下幾乎無能為力,這也提示了NKI編程接口本身還有改進空間。

宅中地 - Facebook 分享 宅中地 - Twitter 分享 宅中地 - Whatsapp 分享 宅中地 - Line 分享
相關內容
Copyright ©2026 | 服務條款 | DMCA | 聯絡我們
宅中地 - 每日更新