近日，來自斯坦福和普林斯頓的研究者發現，DeepSeek-R1已經能生成自定義CUDA內核了，而且還在一眾推理模型中，直接拿下了TOP 1！

緊隨其后，OpenAI o1和Claude 3.5 Sonnet分別排第二和第三。

具體過程，就是給定一個PyTorch程序，讓模型對其優化，然后生成一個包含自定義CUDA內核的PyTorch版本。

在此期間中，模型可以自由決定優化哪些操作，以提高計算效率。

引導模型生成GPU內核，潛力巨大

如今，傳統人工優化內核的方式，在效率上已經不足以應對大量涌現的AI模型架構和硬件平臺。

既然是為了LLM進行優化，那么，我們是否也能夠借助LLM來模擬AI工程師的工作流程，憑借編譯器反饋、硬件規格等豐富的信息，自動編寫出準確且經過優化的內核代碼呢？

為此，研究團隊提出了一種全新的KernelBench框架，用于生成和評估不同AI任務（單個操作、操作序列、端到端架構）的內核，并模擬了工程師迭代優化的過程。

論文地址：https://arxiv.org/abs/2502.10517

GPU的本質，是硬件依賴的。因此，研究者們希望嘗試，看是否能通過以下方式，引導模型生成GPU內核。

首先，向模型提供硬件信息（如內存帶寬、TFLOPS），以針對特定GPU（A100、H100等）進行優化。

然后，要讓模型在上下文中展示代表性的內核優化技巧，例如矩陣乘法中的分塊（tiling）或Flash Attention中的在線softmax。

研究者們發現，只有更強大的模型，會偶爾表現出利用這些優化的能力。

比如，DeepSeek-R1有時會使用特定于硬件的指令（如Tensor Core的wmma），但往往無法正確編譯或使用它們，從而限制了最終性能。

總的來說，研究發現，前沿模型在KernelBench上的開箱即用性較差，OpenAI o1和DeepSeek-R1在不到20%的任務上超過PyTorch Eager基線。

這些模型生成的內核存在大量執行錯誤、功能正確性問題，并且無法進行特定平臺的優化。

具體來說，研究者發現：

1. 對模型而言，編寫功能正確的內核仍然具有挑戰性；
2. 模型通過優化展示了生成高性能內核的潛力；
3. 利用反饋對于減少執行錯誤和發現更快的方案很重要。

當然，KernelBench目前還只是讓GPU加速奔跑的起點，但也是讓整個GPU編程自動化的起始催化劑。

令人興奮的是，現在已經有了許多新的工作，專注于解決KernelBench中涉及的問題。

比如2月12日，英偉達就發出博客文章，探討如何使用DeepSeek-R1進行GPU內核自動生成與推理時scaling。

隨后在2月12日，Meta也發文測試了前沿模型編寫GPU內核方面的性能，他們發現，最佳模型可以在KernelBench上提供平均1.8倍的加速。

Sakana AI更是推出「AI CUDA工程師」，讓AI自己寫代碼優化CUDA內核，速度聲稱比PyTorch原生實現快了10-100倍。

如雨后春筍般出現的研究表明，如今，我們已經進入了AI驅動加速AI的新紀元！

在未來，KernelBench還將持續演進。它不會僅限于當前收集的250個問題，還可以擴展到新的AI任務。與此同時，評測指標fast_p也可以隨著時間的推移進行調整，提高加速門檻，以推動更高效的優化方案

KernelBench：AI內核生成框架

KernelBench是一個開源框架，旨在評估LLM在編寫GPU內核方面的能力。

任務格式

KernelBench包含250個任務，涵蓋了各種AI工作負載，并且易于擴展到新的工作負載。

下圖1展示了KernelBench評估語言模型（LM）生成高性能GPU內核的能力。KernelBench要求語言模型為給定的目標PyTorch模型架構生成優化的CUDA內核，并進行自動化評估。

任務輸入

給定一個AI工作負載，任務的輸入是用PyTorch編寫的參考實現。模仿研究人員的工作流程，PyTorch代碼包含一個繼承自torch.nn.Module ()的名為Model的類，其中標準的__init__和 forward () 函數（以及任何輔助函數）被填充為AI工作負載的PyTorch操作。

AI算法通常在大型張量數據上進行操作。工作負載的最優內核取決于張量的大小和數據類型（如BF16、FP8）。因此，每個任務還包含get_inputs ()和get_init_inputs ()函數，用于指定內核需要處理的精確輸入張量。

任務輸出

給定輸入，LLM需要輸出一個繼承自torch.nn.Module ()的名為ModelNew的新類，其中包含自定義優化。例如，LLM可以在forward ()函數中使用PyTorch的CUDA-C擴展來集成內聯內核調用。

為了成功完成任務，模型需要確定（1）Model類中的哪些操作最能從優化中受益；（2）如何優化這些操作。LLM可以使用任何硬件高效技術（如融合和分塊）或專用指令（如張量核心）以及任何編程庫（如PTX、CUDA、CUTLASS、Triton、ThunderKittens）。

任務選擇

這些任務根據包含的基本操作或PyTorch庫函數的數量分為三個級別。

Level 1包含100個單個基本操作，如卷積、矩陣乘法等AI基礎構建塊。雖然PyTorch調用了經過優化的閉源內核，讓LLM超越基線具有挑戰性，但如果能生成開源內核，將有重要價值。

Level 2包含100個操作序列，如卷積、ReLU和Bias的組合，這些操作可以融合成一個內核以提高性能。

由于基于編譯器的工具（如PyTorch編譯器）在融合方面非常有效，LLM要在這方面超越它們也具有挑戰性。然而，LLM可能會提出更復雜的算法。

Level 3包含50個完整的機器學習架構，如AlexNet和MiniGPT等，這些架構在運行訓練和推理時對內核的性能要求極高。

評估指標

KernelBench引入了一個新的評估指標fast_p，衡量生成的內核中功能正確且加速大于閾值p的任務比例。

通過調整閾值參數p，研究者可以評估不同加速閾值下的內核性能，并捕捉加速分布。

fast_0相當于LLM的正確率，它衡量代碼功能正確的任務比例，而不考慮其速度。在實際評估中，通常以p=1作為起點。

LLM在KernelBench上的表現

研究人員對一系列LLM在KernelBench上進行了評估，結果顯示，目前的LLM在生成正確且優于PyTorch基線速度的內核方面仍有困難。

在一次性基線評估中，LLM生成的內核平均在不到20%的任務中比PyTorch Eager更快。這表明，僅靠簡單提示，LLM很難在性能上超越傳統的PyTorch內核。

LLM生成的內核存在大量的執行錯誤和功能正確性問題，經常由于簡單的編譯器和運行時錯誤而失敗。

執行錯誤包括CUDA/nvcc/Python編譯時錯誤、CUDA內存違規和運行時錯誤等；正確性錯誤則主要表現為輸出張量形狀和值不匹配。

推理模型（o1，R1）生成的錯誤解決方案（<55%）比其他模型（>70%）少。然而，這主要是由于執行失敗的情況較少。在功能正確性方面，所有LLM都面臨類似程度的困難。

在性能方面，模型生成功能正確的內核在多數情況下也未能優于PyTorch基線。

隨著p的提高，模型生成的內核中能達到要求的比例越來越低。在p=1時，在所有KernelBench級別中，不到15%的LLM生成內核優于PyTorch。

推理模型通常在提供加速方面優于其他LLM，但總體仍有不足。

模型生成的內核在不同硬件平臺上的通用性不佳。DeepSeek-R1生成的內核在NVIDIA L40S上實現了36%的加速，而在NVIDIA A10G上則為47%。

這表明LLM在生成特定目標硬件的高效內核方面還存在很大的提升空間。

模型能力分析

測試時利用KernelBench環境反饋

正如上面觀察到的，執行失敗是LM生成的內核中最常見的失敗模式。

KernelBench提供的環境允許收集豐富的信號，包括編譯器錯誤、正確性檢查和運行時性能分析指標，所有這些都可以反饋給LM以幫助它解決內核故障。

為了探索LM如何利用這些反饋，研究團隊評估和比較了兩個基線：第一個令LM為每個KernelBench任務生成多個并行樣本，另一個通過允許LM利用執行反饋逐步改進，依次為每個KernelBench任務生成內核。

重復采樣

KernelBench環境支持對LM生成的內核進行程序化驗證，允許研究團隊收集和評估每個任務的多個LM生成。他們使用fastp@k評估這種重復采樣方法。重復采樣有助于LM發現更多快速且正確的解決方案。

如下圖4所示，隨著k值的增加，在DeepSeek-V3和Llama 3.1 70B的三個級別上，通過高溫度參數重復采樣可以提升fast1的性能。

值得注意的是，在Level 2上，DeepSeek-V3在k=100個樣本時達到了37%的fast1，而在單次提示基線中僅為4%。

通過檢查樣本，我們發現高溫采樣有助于探索解決方案空間，增加了生成具有更好優化的無錯誤內核的機會。然而，如果一個模型解決任務的固有概率非常低，僅僅增加采樣預算的影響有限。

例如，即使嘗試了100個樣本，DeepSeek-V3也從未能夠為Level 1中的一組34個卷積變體生成任何正確的解決方案。

生成結果的迭代優化

KernelBench環境非常適合收集編譯器反饋、執行錯誤和使用PyTorch分析器等工具進行的時間分析，作為真實信號（ground-truth signals）。

研究人員研究了利用這些反饋是否能幫助語言模型（LMs）迭代優化其生成結果。

下圖5顯示，KernelBench框架使模型能夠在迭代優化過程中接收并利用反饋。這些真實信號包括NVCC編譯器錯誤信息、執行統計數據（例如正確性檢查和掛鐘時間），以及PyTorch分析器（操作時間分解）。

他們在多輪過程中為模型提供每次生成的反饋：在初始生成后，向模型提供其之前的生成結果G，以及當前生成對應的編譯器/執行反饋E和/或分析器輸出P。

然后將每次生成及其后續反饋定義為一輪（turn），并在N輪內運行這一迭代優化過程。利用執行反饋有助于減少錯誤，并隨時間提升整體加速效果。

研究人員在下表2中檢查了第N=10輪時的fast1行為，發現迭代優化在不同模型和KernelBench的各個級別上均持續提升了性能。

DeepSeek-R1在Level 2上的改進最為顯著，其中執行反饋E和分析器反饋P的組合將fast1從36%提升至72%（如下圖6所示）。

此外，通過分析迭代優化軌跡，他們發現模型在執行反饋E的幫助下能更有效地自我糾正，尤其是在修復與執行錯誤相關的問題上。

DeepSeek-R1在Level 1和Level 2上，經過10輪優化后，能在超過90%的任務中生成功能正確的內核（下表9）。

然而，剩余的錯誤內核幾乎總是由于功能不正確而失敗，這可能是因為正確性反饋的顆粒度不如執行失敗信息細致。

比較重復采樣與迭代優化

在上表2中，研究人員比較了在固定10次推理調用預算下重復采樣和迭代優化的效果。兩種方法相較于單次基線（one-shot baseline）均取得了顯著改進，其中迭代優化在6個案例中的5個中表現更優。

然而，他們最終發現，測試時方法的效果本質上依賴于基礎模型的質量。

例如，在重復采樣中，DeepSeek-V3在所有三個級別上始終優于Llama-3.1 70B。類似地，在迭代優化中，DeepSeek-R1通過反饋E和P持續改進，而DeepSeek-V3和Llama-3.1 70B并非總能從這些信息中獲益。

提供硬件知識生成硬件高效內核

顯然，語言模型在生成硬件高效內核方面表現有限。

這可能是由于訓練數據中內核代碼的稀缺性，以及最佳內核可能需要根據硬件平臺的特定屬性而變化。

在本案例研究中，研究團隊探索了提供以下內容的效果：（1）提供內核工程最佳實踐的示例，并將其置于（語言模型的）上下文之中；（2）提供詳細的硬件規格說明，并將其置于（語言模型的）上下文之中。

硬件感知的上下文示例

編寫良好的內核通常使用融合（fusion）、分塊（tiling）、重計算（recompute）和異步（asynchrony）等技術來最大化性能。

具體來說，研究人員納入了三個上下文示例：使用操作融合的GeLU、使用分塊的矩陣乘法，以及展示共享內存I/O管理的最小Flash-Attention內核。

結果則顯示，上下文示例降低了語言模型的整體fast1分數，因為模型嘗試了更激進的優化策略，但導致更多執行失敗。與上面基線生成的代碼相比，OpenAI o1在使用少樣本示例時生成的代碼平均長度增加了25%。

然而，在正確的解決方案中，語言模型應用了一些有趣的優化：他們發現，在KernelBench Level 1的77%的GEMM變體中，o1應用了分塊并提升了速度，優于單次基線。在Level 2，o1在11個問題上應用了激進的共享內存I/O管理，并能夠超越PyTorch Eager。

指定硬件信息

正如上面討論過的，內核性能因硬件平臺而異。

例如，FlashAttention-2從NVIDIA A100遷移到H100 GPU時硬件利用率下降了47%。FlashAttention-3是一個完全不同的算法，專為H100編寫。

在本研究中，研究團隊探討語言模型是否能利用上下文中的以下信息生成改進的內核：（1）硬件規格，例如GPU類型（H100、A100等）、內存大小、帶寬、TFLOPS；（2）硬件知識（例如線程、線程束、線程塊、流多處理器的定義）。

結果顯示，模型很少生成針對底層硬件優化的內核，這表明未來模型仍有改進空間。

某些新一代GPU（例如H100）引入了與前代不同的新硬件單元和指令。提供硬件信息對Llama 3.1 70B或DeepSeek-V3的輸出影響不大。

有趣的是，他們發現OpenAI o1和DeepSeek-R1生成的部分內核使用了特定于硬件的指令和優化。

R1在大約50%的Level 1矩陣乘法問題中嘗試生成warp矩陣乘加（wmma）指令（下圖10），盡管大多數未能編譯。

在功能正確的生成中，R1和o1在每個級別產生了1-3個異常值，比Level 4基線快2倍以上。

總體而言，研究團隊發現，與提供硬件信息相比，語言模型通過少樣本示例調整策略時表現更好。

結論

研究人員在本論文中提出了KernelBench，一個為語言模型驅動的內核優化奠定基礎的框架；他們評估了多種模型和方法，分析了它們的優勢和局限性，并提供了改進機會的見解。

總的來說，盡管大多數基準測試最終會達到飽和，但KernelBench被設計為隨著新的AI工作負載的出現而動態演進。

他們的fastp指標可以隨時間調整，以測量相對于日益先進的基線（即超出工作中使用的PyTorch基線）的加速閾值（p）。

由于PyTorch具備跨硬件平臺兼容性，KernelBench中基于PyTorch的任務可以在每個新硬件平臺發布時進行評估。最后，與許多基準測試不同，在KernelBench上的成功直接映射到生產價值和現實世界的影響（降低成本并大規模減少能耗）。

這些特性確保了KernelBench在不斷演變的AI領域中將保持其價值。

下一步工作

研究團隊表示在當前可用模型下，KernelBench仍有顯著的改進空間。

首先，未來的工作可以探索開發先進的微調和推理技術，包括智能體工作流（agentic workflows）。由于CUDA是一種低資源語言，未來工作開源更多高質量數據將具有重要價值。

其次，在他們的實驗中，語言模型生成的是原始CUDA代碼。然而，未來的工作可以探索使用其他編程抽象（例如ThunderKittens、CUTLASS、Triton等）生成代碼是否能簡化生成問題，例如使語言模型更容易利用張量核心指令。

最后，研究團隊的評估至今僅限于GPU，未來的工作可以擴展到其他硬件加速器。

作者介紹

Anne Ouyang

Anne Ouyang目前是斯坦福大學計算機科學（CS）博士生，在Scaling Intelligence Lab（規模化智能實驗室）進行研究。

她的研究興趣主要集中在可擴展的自我改進機器學習系統，同時也廣泛關注實證機器學習（empirical ML）和性能工程（performance engineering）。

她曾獲得了MIT學士和碩士學位，并曾在NVIDIA cuDNN團隊工作，負責編寫CUDA內核，用于加速GPU上的深度學習工作負載。

Simon Guo

Simon Guo是斯坦福大學計算機科學專業的一年級博士生，目前正在擴展智能實驗室（Scaling Intelligence Lab）跟隨Azalia Mirhoseini教授進行輪轉研究。

他曾獲得了UC伯克利電氣工程和計算機科學學士學位。他的研究興趣在計算機系統和機器學習。

最近，他在Cohere從事語言模型預訓練工作。在此之前，他曾在蘋果公司設計GPU，在Anyscale開發分布式系統，并在NVIDIA DRIVE部門從事自動駕駛汽車的開發工作。