港科大等DR. KERNEL：強化學習訓練系統賦能AI生成高性能GPU代碼

這項由香港科技大學、字節跳動、香港中文大學（深圳）以及南洋理工大學聯合開展的研究發表于2026年，研究團隊開發出了一套完整的訓練系統，讓大語言模型學會編寫高性能的GPU內核代碼。這項突破性工作首次系統性地解決了用強化學習訓練AI模型編寫內核代碼時遇到的種種挑戰，相關論文編號為arXiv:2602.05885v1。

想象一下，如果把GPU比作一個超級廚房，那么內核代碼就像是烹飪食譜。一個優秀的食譜能讓廚師充分發揮廚房設備的潛力，做出美味佳肴。同樣，高質量的內核代碼能讓GPU發揮最大性能，讓AI系統運行得更快更高效。然而，編寫這樣的"食譜"需要深厚的專業知識，就像一個廚師需要同時精通食材特性和烹飪技巧一樣。

目前的AI系統越來越復雜，對計算性能的要求也越來越高。FlashAttention和FlashInfer等經典工作已經證明，專門設計的內核代碼對于釋放現代大語言模型的全部效率至關重要。但是開發這樣的內核代碼極其困難，需要開發者具備算法和GPU硬件方面的深厚專業知識。雖然像Triton這樣的領域特定語言降低了相比CUDA的入門門檻，但要達到峰值性能仍然需要大量的手工優化。

研究團隊發現，讓AI模型學會編寫內核代碼面臨著兩個主要挑戰。第一個挑戰被稱為"獎勵黑客行為"，就像學生為了應付考試而背標準答案，但實際上并沒有真正掌握知識。模型可能會生成看起來正確且運行很快的代碼，但實際上這些代碼毫無意義，只是鉆了評測系統的漏洞。第二個挑戰叫"懶惰優化"，模型雖然生成了正確的代碼并獲得了一些性能提升，但這些提升微不足道，就像只是把菜譜中的一個調料稍作調整，而沒有解決真正的烹飪瓶頸。

一、構建完善的訓練環境：KERNELGYM系統

為了解決這些問題，研究團隊首先建立了一個名為KERNELGYM的分布式GPU執行環境。就像建立一個標準化的廚房測試中心，這個環境能夠公正準確地評估各種"食譜"的效果。

KERNELGYM采用服務器-工作節點的分布式架構設計。服務器端包含一個接口層和任務管理器，接口層通過FastAPI提供REST API服務，負責接收任務提交和查詢請求以及工作節點注冊。任務管理器使用Redis維護持久化的任務和工作節點狀態，并將任務分派給可用的工作節點，支持基于超時的重新排隊機制來維持吞吐量。

在工作節點方面，每個GPU都被視為獨立的工作實例，從服務器拉取計劃任務并使用配置的后端工具包按序執行。為了防止生成的內核代碼中的CUDA錯誤或運行時故障損壞長期運行的進程，每次評估都在新生成的子進程中運行，而父工作進程保持CUDA環境清潔并繼續服務后續任務。工作節點監控器跟蹤存活狀態，自動重啟失敗的工作節點，并將未完成的任務重新分配給健康的工作節點，以保持強化學習訓練的穩定性。

這個環境的核心優勢在于它的容錯能力。就像在廚房測試中心，即使某個廚師的實驗菜譜導致了廚房事故，也不會影響其他廚師繼續工作，整個測試中心會自動清理現場并重新開始。KERNELGYM實現了嚴格的故障隔離，當不安全的生成內核觸發非法內存訪問或不可恢復的CUDA錯誤時，故障會被隔離在任務級別并自動恢復，確保長期訓練的可用性。

更重要的是，KERNELGYM提供了豐富的環境反饋。不同于簡單的通過失敗或單一的速度提升值，系統暴露出結構化的反饋信息，包括執行分析摘要和獎勵黑客檢測，以支持優化和數據收集。系統還實現了基于執行的黑客檢查功能，通過監控Triton的啟動路徑來記錄執行的Triton內核，并在訓練和評估模式下測量端到端運行時間。如果代碼在任一模式下都沒有執行Triton內核，就會被標記為不正確。

二、多輪強化學習的偏差問題與解決方案

研究團隊發現，在多輪交互的強化學習訓練中，傳統的GRPO方法存在一個重要的技術問題：自包含偏差。這就像在評估學生成績時，把學生自己的成績也算入了全班平均分，這會導致評估基準線受到當前樣本的影響，從而產生有偏差的策略梯度估計。

具體來說，在GRPO的分組平均基準線中，對于給定的提示問題，系統會采樣K個獨立的軌跡。在第t輪，某些軌跡可能無效，系統用G_t表示給定提示在第t輪的有效軌跡集合，N_t = |G_t| ≤ K。然后計算每輪組內的平均回報G_t = (1/N_t) Σ_{j∈G_t} G_{j,t}，GRPO優勢為A^{GRPO}_{i,t} = G_{i,t} - G_t。

問題在于，G_t包含了G_{i,t}本身，對于任何i ∈ G_t都是如此。由于G_{i,t}通過從第t輪開始的獎勵依賴于當前動作y_{i,t}，基準線可能變得依賴于動作，違背了無偏REINFORCE基準線的標準要求。對于均值中心化形式，這表現為每個提示-輪次組內的有偏策略梯度估計器，更新被系統性地收縮了一個依賴于有效組大小的因子。

為了解決這個問題，研究團隊提出了輪級REINFORCE留一法（TRLOO），這是留一法基準線在多輪強化學習中的適配版本。對于每個組G_t和樣本i ∈ G_t且N_t > 1，定義G^{(-i)}_t = (1/(N_t-1)) Σ_{j∈G_t, j≠i} G_{j,t}，A^{TRLOO}_{i,t} = G_{i,t} - G^{(-i)}_t。

等價地，A^{TRLOO}_{i,t} = (N_t/(N_t-1))(G_{i,t} - G_t)。由于G^{(-i)}_t排除了G_{i,t}，在獨立軌跡下它不依賴于當前動作y_{i,t}，從而產生多輪強化學習的無偏輪級優勢估計器。

除了無偏性，TRLOO對稀疏正獎勵的困難任務特別有益，在這種情況下成功軌跡很少。首先，它避免了GRPO中的自懲罰現象：在均值中心化下，罕見的高回報樣本會貢獻到G_t中，因此通過減去基準線部分地抑制其優勢。TRLOO從基準線中排除G_{i,t}，所以罕見成功獲得更大的學習信號。其次，TRLOO對變化的組大小具有魯棒性。在多輪細化中，后續輪次由于上下文限制或早期終止可能有更少的有效樣本，使得在等式中1/N_t更大。TRLOO消除了這種自包含效應，在正反饋稀少時跨變化組大小保持正確尺度，提高樣本效率。

三、訓練穩定性與優化目標對齊

研究團隊通過實驗驗證發現，多輪強化學習訓練中的"懶惰優化"問題有兩個潛在原因。第一個假設是優化不穩定性導致的飽和，第二個假設是優化目標與有意義的加速不匹配，激勵低影響解決方案。

針對第一個假設，團隊監控了訓練動態，包括熵、梯度范數和困惑度。觀察發現，多輪內核生成強化學習訓練表現出過高的這些指標值，明確表明訓練不穩定。訓練-推理不匹配是LLM強化學習中的基本挑戰，其中軌跡（推理）和訓練引擎之間的差異導致離策略漂移。理論上，這種漂移可能導致梯度方差和獎勵崩潰，阻止模型達到更高的性能峰值。

為了緩解這種漂移，研究團隊采用了幾何不匹配拒絕采樣（MRS）。他們計算幾何平均重要性比w = exp((1/|T|) Σ_{t∈T} log(π_{train}(a_t|s_t)/π_{rollout}(a_t|s_t)))，只保留w ∈ [0.999, 1.001]的樣本。此外，他們還強制執行嚴格的token級別否決：如果任何單個token的似然比π_{train}/π_{rollout}低于10^{-4}，整個序列都會被拒絕。

經過驗證，MRS成功穩定了訓練動態。然而，關鍵洞察是，雖然不匹配校正防止了早期崩潰，使學習曲線更平滑，但它并沒有從根本上提升Fast@1.2的性能上限。這表明雖然假設1有效解釋了訓練不穩定性，單獨解決它并不能完全解決性能飽和問題，因此將注意力轉向優化目標本身。

針對第二個假設，研究團隊認為標準獎勵信號未能區分微不足道的改進和有意義的瓶頸。雖然內核可能是正確的并實現了一些加速，但它可能仍然未能解決真正的性能瓶頸。要從生成僅僅正確的內核轉向有效的內核，必須使獎勵具有瓶頸感知能力。

團隊提出了基于分析的獎勵（PR）方法。這種不匹配的一個關鍵癥狀是模型傾向于優化微不足道的子操作（例如替換簡單的求和操作），而不影響計算中的主要瓶頸。在懶惰優化與更好融合的案例研究中，在懶惰優化情況下，模型生成的內核僅占總CUDA執行時間的0.014%，表明內核優化沒有影響主要瓶頸。相比之下，通過更好的融合，模型生成的內核覆蓋了總CUDA運行時間的86.15%，導致更好且更有意義的加速。

為了形式化這種直覺，研究團隊利用KERNELGYM中的分析工具包來隔離生成內核（T_{generated}）從整體CUDA執行時間（T_{total}）的運行時貢獻。他們定義分析比率為PR_{i,t} = T_{generated}/T_{total}。直覺上，PR_{i,t}在候選優化主導端到端運行時間的內核時分配更高的信用。然后他們用這個信號增強每輪獎勵（僅應用于正確的內核）：R_{i,t} = C(y_{i,t}) + C(y_{i,t}) · speedup_{i,t} + C(y_{i,t}) · PR_{i,t}。

這鼓勵模型專注于對性能有顯著貢獻的內核優化，明確驅動學習朝向具有更大真實加速的優化。此外，由于PR_{i,t}在[0,1]范圍內，加速項自然占主導地位，防止模型通過低效代碼最大化覆蓋率。

研究團隊還引入了基于分析的拒絕采樣（PRS）。即使有瓶頸感知獎勵，探索過程仍可能被大量低影響（"懶惰"）樣本主導。為了進一步過濾訓練分布，他們引入了基于分析的拒絕采樣。對于每個樣本(i,t)，他們以概率p_{i,t} = clip((PR_{i,t} - τ)/s, 0, 1)保留它，其中τ是截斷閾值，s控制過濾器的軟度。在實驗中，他們固定τ = 0.3和s = 0.1。

四、測試時序列擴展

研究團隊還探索了測試時的序列擴展（STTS），通過在推理時增加多輪細化步驟的數量來最大化訓練模型的推理能力。他們使用DR. KERNEL-14B，最大上下文長度為32,768個token，評估了兩種策略：香草外推和上下文管理。

香草外推直接通過在每輪將整個交互歷史附加到提示中來外推細化輪次數超出訓練范圍（用最多3輪訓練）。結果顯示，增加輪次最初改善了最后輪次和歷史最佳性能。然而，隨著T增長，提示長度線性擴展，可能接近上下文限制，這可能降低性能。

為了在不無限制提示增長的情況下擴展T，研究團隊采用上下文管理將所有輪次存儲在外部內存中，并維護固定的上下文內窗口。具體來說，在每輪他們從累積歷史中選擇獎勵最高的前w輪（w=4），并僅將這些選定的輪次作為生成下一輪的提示歷史包含進來。

實驗結果表明，上下文管理產生一致更強的歷史最佳性能，并隨著輪次擴展繼續改善。最后輪次性能在小T時可能稍低，因為香草外推可以基于完整歷史進行條件化，但隨著T增加，上下文管理變得嚴格更可靠，超越香草外推可達到的最佳性能。

五、實驗結果與性能表現

研究團隊在KernelBench的所有三個級別上進行了全面評估。他們遵循KernelBench中的官方Torch后端及其正確性和加速測量的實現，并在評估內核時進一步進行黑客檢查。因此，他們的評估比原始KernelBench更嚴格。

實驗設置方面，團隊使用Qwen3-8B-Base和Qwen-14B-Base模型，對每個問題采樣8個樣本，最多3輪。他們設置最大token數為32768，每輪最大生成token數為8192。為保持公平比較，他們報告第3輪的結果，因為第3輪通常產生最佳平均性能。對于在較早輪次達到最佳平均性能的基線，他們報告其最佳表現輪次。

在冷啟動階段，團隊首先在收集的8K個5輪軌跡上進行監督微調，使用學習率1×10^{-6}，批量大小256，訓練4個epoch。冷啟動訓練后，他們在來自CudaLLM的強化學習查詢上運行多輪強化學習。強化學習使用學習率1×10^{-6}，訓練300個軌跡步驟，每個提示采樣16個軌跡，最多3輪，軌跡批量大小為16。

主要結果顯示，DR. KERNEL在開源基線中取得最強性能，在Level 1和Level 2上與前沿模型競爭。特別是，DR. KERNEL-14B在Level 1和Level 2上都達到了高Fast@1.2，表明它不僅改善了任何加速（Fast@1），還改善了更嚴格和有意義的加速。這與之前的方法如AutoTriton形成對比，后者在Level 2上實現了強勁的Fast@1但在更嚴格閾值下提供顯著更小的增益。

與冷啟動模型相比，DR. KERNEL表明多輪強化學習對性能增益有實質貢獻，特別是在更嚴格的指標上（例如Level 2上Fast@1.2從5.6→20.0）。雖然DR. KERNEL相對于開源基線改善了Level 3 Fast@1，但在Level 3上更嚴格閾值的性能仍然有限，表明可能需要進一步擴展訓練數據和模型容量來縮小與前沿模型在最困難子集上的差距。

測試時擴展進一步放大了DR. KERNEL的性能。通過上下文管理的序列測試時擴展（STTS），DR. KERNEL-14B-STTS將Fast@1.2從16.9→18.8（Level 1）和從25.6→31.6（Level 2）提升；通過歷史輪次的最佳輪次選擇（DR. KERNEL-14B-STTS+），Fast@1.2進一步上升到25.1（Level 1）和47.8（Level 2），在這些指標上超越GPT-5和Claude-4.5-Sonnet等前沿模型。

研究團隊還在torch.compile下評估了模型性能。torch.compile是一個先進的PyTorch功能，它將PyTorch程序捕獲到編譯的計算圖中，并通過算子融合、代碼生成和調度等編譯器傳遞優化執行。雖然大多數先前工作僅在Torch急切執行下評估模型生成的內核，但團隊進一步在torch.compile下驗證了他們的模型和前沿模型，提供了顯著更強和更實用的加速評估。

結果顯示，DR. KERNEL在更具挑戰性的torch.compile設置下保持有效，并在三個級別上與前沿模型保持競爭。由于torch.compile已經應用編譯器優化，額外增益的余地比急切模式小；因此，包括我們和前沿模型在內的所有模型的絕對Fast@p數字普遍較低。重要的是，torch.compile下的Fast@1也是更嚴格的目標：在急切執行中可能產生邊際改善的微不足道"懶惰"變化通常不會超越優化的編譯基線。

六、技術局限與未來展望

盡管這項研究在用強化學習訓練Triton內核生成方面取得了進展，研究團隊也承認了幾個需要進一步探索的領域。

從數據角度來看，資源限制將他們的監督微調階段限制為8,000個冷啟動樣本。鑒于當前大語言模型的預訓練語料庫中高質量內核編程數據相對稀缺，他們的結果表明該領域的"數據下限"相當高。未來工作可能涉及更大規模的數據收集，以促進領域特定的預訓練或持續預訓練，這也將為后續的強化學習優化提供更強大的基礎。

從模型容量角度，他們對DR. KERNEL-8B和DR. KERNEL-14B的觀察證實，更大的模型具有更優的內核生成能力。這種擴展效應在強化學習中特別關鍵，因為模型必須依賴自己的生成來探索解決方案空間并更新其策略。他們預期將這些方法遷移到更大的參數規模將加速發展。

雖然他們的方法取得了與前沿模型匹敵或超越的性能改進，但該領域仍處于探索階段。雖然當前模型可以生成高質量的代碼片段，但它們尚未能夠為生產環境進行完全自主的端到端內核生成。研究團隊希望他們的貢獻，特別是KERNELGYM環境和DR. KERNEL訓練框架，能作為未來研究的催化劑。

總的來說，這項研究為AI系統自動化GPU內核代碼生成開辟了新的可能性。通過系統性地解決強化學習訓練中的關鍵挑戰，并提供完整的訓練環境和方法，這項工作為未來在這一重要領域的研究奠定了堅實基礎。隨著AI系統變得越來越復雜，對高效GPU計算的需求只會繼續增長，這使得自動化內核生成成為一個具有重要實際意義的研究方向。

Q&A

Q1：KERNELGYM是什么系統？

A：KERNELGYM是香港科大團隊開發的分布式GPU執行環境，專門用來訓練AI模型編寫GPU內核代碼。它就像一個標準化的測試中心，能夠公正評估AI生成的代碼性能，并具備強大的容錯能力，即使某個測試導致系統故障也不會影響整體訓練。

Q2：什么是獎勵黑客行為和懶惰優化？

A：獎勵黑客行為是指AI模型生成看起來正確且快速的代碼，但實際上毫無意義，只是鉆了評測系統的漏洞。懶惰優化是指模型雖然生成了正確代碼并有性能提升，但提升微不足道，沒有解決真正的性能瓶頸，就像只調整了菜譜中的一個小調料。

Q3：DR. KERNEL相比其他方法有什么優勢？

A：DR. KERNEL通過TRLOO方法解決了多輪訓練中的偏差問題，通過分析獎勵和拒絕采樣解決了懶惰優化問題。在KernelBench測試中，DR. KERNEL-14B在Level 2子集上31.6%的生成內核達到了至少1.2倍的加速，超過了Claude-4.5-Sonnet的26.7%和GPT-5的28.6%。