多智能體接管 TPU 底層調校:MaxKernel 讓 LLM 寫出媲美專家的加速核
huggingface.co/papers · 2026-09-07
Google DeepMind 團隊於 2026 年 9 月發布 MaxKernel,這是一套讓大型語言模型結合即時編譯器回饋、自動產生TPU 底層運算核心(kernel)的多智能體系統。團隊在自建的 JaxBench 基準上測試,產生的 kernel 效能可媲美甚至超越人類專家手寫版本。論文以 arXiv:2609.04523 發布,程式碼同步公開於 GitHub。
背景:為什麼 TPU kernel 要手工調校
TPU 是 Google 自研的 AI 加速晶片,實際運算效能高度仰賴針對特定操作(如矩陣乘法、注意力機制)撰寫的低階運算核心。這些 kernel 通常以 Pallas 一類的領域特定語言撰寫,工程師需手動決定資料分塊大小、記憶體佈局與流水線排程,才能吃滿晶片的算力與頻寬。每當硬體世代更新或模型架構改變,過去調好的 kernel 往往需要重新優化,而具備這種底層調校能力的工程師相對稀少,形成產能瓶頸。
核心方法:三種智能體協作模式
MaxKernel 把 kernel 開發拆解成規劃、實作、驗證、自動調參、剖析等子任務,分別交給專責的子智能體處理,並將編譯器與效能分析工具的回饋即時餵回給智能體修正程式碼。系統提供三種可切換的運作模式:人機協作(HITL)讓工程師在關鍵節點審核與介入,自主迴圈模式讓智能體依效能指標持續規劃、編譯、測試並反覆優化,圖搜尋模式則把設計空間形式化為搜尋圖,以平行或束搜尋探索更廣的候選解。
為了避免智能體直接抄襲既有寫法,系統採用 RAG 技術檢索的內容僅限硬體規格文件與編譯器手冊,不包含任何人類寫好的 kernel 範例,迫使模型從底層原理推導寫法而非模式匹配。
影響範圍:效能可比甚至超越人類基準
在涵蓋 50 個多樣任務的 JaxBench 上,平行搜尋模式對照專家基準達到幾何平均1.58 倍加速,束搜尋則為 1.49 倍,50 個任務全數產出可編譯且結果正確的 kernel。換算到實際生產環境的 kernel 集合,平行搜尋模式達到 2.32 倍幾何平均加速,略高於人類工程師手寫版本的 2.02 倍。
針對真實 SOTA 模型的個別 kernel,系統在 DeepSeek-V4 的某個運算核心上達到7.85 倍加速,在 Qwen3-Next 的訓練流程上帶來 4.70 倍加速,顯示這套流程在超出訓練分佈的新架構上仍具備泛化能力。
| 評測情境 | 加速倍數(相對基準) |
|---|---|
| JaxBench(50 任務)· 平行搜尋 | 1.58× |
| JaxBench(50 任務)· 束搜尋 | 1.49× |
| 生產環境 kernel · 平行搜尋 | 2.32×(人類基準 2.02×) |
| DeepSeek-V4 單一 kernel | 7.85× |
| Qwen3-Next 訓練流程 | 4.70× |
別急著丟掉 Dropout:Cerebras 重新驗證「層級丟棄」對 LLM 訓練與推論的雙重效益
arxiv.org · 2026-09-07
Cerebras 團隊發表研究,重新檢視在大型語言模型預訓練中普遍遭到棄用的dropout 技巧,證明只要調整成「整層丟棄」並搭配正確設定,不但不會拖累收斂,反而能在同等算力下降低驗證損失。論文已被 ICML 2026 接受,以 arXiv:2609.05275 公開,實驗全數在 Cerebras CS-3 系統上完成。
背景:dropout 為何從 LLM 訓練中消失
Dropout 是類神經網路的經典正則化手法:訓練時以一定機率隨機捨棄部分神經元的輸出,迫使網路不依賴單一路徑,藉此抑制過擬合。但隨著模型規模進入數十億參數,業界逐漸觀察到 dropout 會拖慢收斂速度、在等量訓練步數下墊高損失,因此主流 LLM 預訓練配方多半直接關閉 dropout。「層級丟棄」(layer dropout)是更粗粒度的變體,不是丟棄個別神經元,而是以一定機率讓整條 transformer 層在該次前向傳播中被跳過。
核心方法:大規模系統性重新掃描超參數空間
作者在 271M 到 8.2B 參數的模型上,搭配最多 160B tokens 的訓練資料,總計執行超過2,400 組實驗,系統性掃描層級丟棄率在深度方向的分布方式、丟棄率隨訓練進程調整的時間表,以及它與 optimizer 超參數之間的交互作用。研究發現先前認為 dropout 不利於 LLM 訓練的結論,很大程度來自沒有為 layer dropout 重新校準這些設定,一旦用正確的層分布與排程,結果會反轉。
影響範圍:訓練與推論的雙重收益
在最佳配置下,層級丟棄能在相同驗證損失下省下最多 25% 的訓練 FLOPs,相當於用更少算力達到同等模型品質。
訓練時隨機跳過層的機制,也讓模型在推論階段可以真正拿掉部分層而只造成極小的準確率損失,論文測得的加速幅度最高達 1.5 倍。這代表同一套訓練技巧同時作用於訓練成本與服務成本兩端,而非常見的「訓練換推論」互斥取捨。
- 系統性重新掃描層分布、時間表與 optimizer 超參數,修正過去對 layer dropout 的負面結論
- 等量驗證損失下最多節省 25% 訓練 FLOPs
- 推論階段可省略部分層,帶來最高 1.5 倍加速且準確率損失極小
當負載平衡「用力過猛」:MESA 修正 MoE 專家剪枝在過度分散路由下的失準
arxiv.org · 2026-09-07
IBM 研究團隊(Berkcan Kapusuzoglu 等七人)發表論文,指出對混合專家(MoE)模型常見的「依 router 機率剪枝專家」做法,在訓練時採用高強度負載平衡的模型上會完全失準,並提出 MESA 修正這個問題。論文於 2026 年 9 月 3 日提交,以 arXiv:2609.04453 公開,長達 22 頁。
背景:MoE 的路由、負載平衡與專家剪枝
混合專家(MoE)模型把原本單一稠密的前饋層,換成多個平行的「專家」子網路,每個 token 經過一個輕量的 router,依機率挑出其中少數幾個專家來實際運算,讓模型在推論成本不變的前提下大幅擴充容量。但若訓練時不同專家的使用率差距過大,會造成部分專家過熱、部分閒置,拖累訓練穩定性,因此業界普遍加入「負載平衡損失」逼迫 token 更均勻地分散到各專家。部署階段常見的壓縮手段是專家剪枝:把 router 給每個專家的平均機率當作重要性分數,直接砍掉分數低的專家來縮減模型體積與服務成本。
核心發現:過度分散路由讓重要性訊號失真
作者發現,當訓練時的負載平衡強度夠高,token 幾乎均勻分散到所有專家,即進入「過度分散路由」狀態,router 機率會趨於一致,不再能反映哪個專家真正重要。在 gpt-oss-20B 上驗證這個現象:選出整體 perplexity 最低的剪枝組合,數學推理能力反而是所有組合裡最差的,相反地,perplexity 最高的組合卻保住了最好的數學推理能力。
作為對照,在採用標準路由的 Mixtral-8x7B-Instruct 上,perplexity 與準確率仍同步下降,沒有出現這種背離,顯示問題只發生在負載平衡被推得很極端的模型上。
MESA 方法與影響範圍
針對這個問題,作者提出 MESA(Minimax Expert Score Allocation):不以單一整體平均準確率為目標,而是採用minimax 策略——反覆找出目前退化最嚴重的領域,調高該領域所依賴專家的重要性分數,直到剪枝後各領域的最差表現被壓到最低。
在 gpt-oss-20B、gpt-oss-120B、Gemma-4-26B-A4B、OLMoE-1B-7B 四款模型上測試,25% 剪枝比例下 MESA 造成的最差領域退化幅度最小;在 11 個 benchmark 中,MESA 於其中 7 個上超越既有的 activation-aware 剪枝基準;在 GPQA 上,不同評分方法之間出現達 18 分的差距。