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

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

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

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

引導(dǎo)模型生成GPU內(nèi)核，潛力巨大

如今，傳統(tǒng)人工優(yōu)化內(nèi)核的方式，在效率上已經(jīng)不足以應(yīng)對大量涌現(xiàn)的AI模型架構(gòu)和硬件平臺。

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

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

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

GPU的本質(zhì)，是硬件依賴的。因此，研究者們希望嘗試，看是否能通過以下方式，引導(dǎo)模型生成GPU內(nèi)核。

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

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

研究者們發(fā)現(xiàn)，只有更強大的模型，會偶爾表現(xiàn)出利用這些優(yōu)化的能力。

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

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

這些模型生成的內(nèi)核存在大量執(zhí)行錯誤、功能正確性問題，并且無法進行特定平臺的優(yōu)化。

具體來說，研究者發(fā)現(xiàn)：

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

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

令人興奮的是，現(xiàn)在已經(jīng)有了許多新的工作，專注于解決KernelBench中涉及的問題。

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

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

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

如雨后春筍般出現(xiàn)的研究表明，如今，我們已經(jīng)進入了AI驅(qū)動加速AI的新紀(jì)元！

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

KernelBench：AI內(nèi)核生成框架

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

任務(wù)格式

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

下圖1展示了KernelBench評估語言模型（LM）生成高性能GPU內(nèi)核的能力。KernelBench要求語言模型為給定的目標(biāo)PyTorch模型架構(gòu)生成優(yōu)化的CUDA內(nèi)核，并進行自動化評估。

任務(wù)輸入

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

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

任務(wù)輸出

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

為了成功完成任務(wù)，模型需要確定（1）Model類中的哪些操作最能從優(yōu)化中受益；（2）如何優(yōu)化這些操作。LLM可以使用任何硬件高效技術(shù)（如融合和分塊）或?qū)Ｓ弥噶睿ㄈ鐝埩亢诵模┮约叭魏尉幊處欤ㄈ鏟TX、CUDA、CUTLASS、Triton、ThunderKittens）。

任務(wù)選擇

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

Level 1包含100個單個基本操作，如卷積、矩陣乘法等AI基礎(chǔ)構(gòu)建塊。雖然PyTorch調(diào)用了經(jīng)過優(yōu)化的閉源內(nèi)核，讓LLM超越基線具有挑戰(zhàn)性，但如果能生成開源內(nèi)核，將有重要價值。

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

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

Level 3包含50個完整的機器學(xué)習(xí)架構(gòu)，如AlexNet和MiniGPT等，這些架構(gòu)在運行訓(xùn)練和推理時對內(nèi)核的性能要求極高。

評估指標(biāo)

KernelBench引入了一個新的評估指標(biāo)fast_p，衡量生成的內(nèi)核中功能正確且加速大于閾值p的任務(wù)比例。

通過調(diào)整閾值參數(shù)p，研究者可以評估不同加速閾值下的內(nèi)核性能，并捕捉加速分布。

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

LLM在KernelBench上的表現(xiàn)

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

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

LLM生成的內(nèi)核存在大量的執(zhí)行錯誤和功能正確性問題，經(jīng)常由于簡單的編譯器和運行時錯誤而失敗。

執(zhí)行錯誤包括CUDA/nvcc/Python編譯時錯誤、CUDA內(nèi)存違規(guī)和運行時錯誤等；正確性錯誤則主要表現(xiàn)為輸出張量形狀和值不匹配。

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

在性能方面，模型生成功能正確的內(nèi)核在多數(shù)情況下也未能優(yōu)于PyTorch基線。

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

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

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

這表明LLM在生成特定目標(biāo)硬件的高效內(nèi)核方面還存在很大的提升空間。

模型能力分析

測試時利用KernelBench環(huán)境反饋

正如上面觀察到的，執(zhí)行失敗是LM生成的內(nèi)核中最常見的失敗模式。

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

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

重復(fù)采樣

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

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

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

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

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

生成結(jié)果的迭代優(yōu)化

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

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

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

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

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

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

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

此外，通過分析迭代優(yōu)化軌跡，他們發(fā)現(xiàn)模型在執(zhí)行反饋E的幫助下能更有效地自我糾正，尤其是在修復(fù)與執(zhí)行錯誤相關(guān)的問題上。

DeepSeek-R1在Level 1和Level 2上，經(jīng)過10輪優(yōu)化后，能在超過90%的任務(wù)中生成功能正確的內(nèi)核（下表9）。

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

比較重復(fù)采樣與迭代優(yōu)化

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

然而，他們最終發(fā)現(xiàn)，測試時方法的效果本質(zhì)上依賴于基礎(chǔ)模型的質(zhì)量。

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

提供硬件知識生成硬件高效內(nèi)核

顯然，語言模型在生成硬件高效內(nèi)核方面表現(xiàn)有限。

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

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

硬件感知的上下文示例

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

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

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

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

指定硬件信息

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

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

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

結(jié)果顯示，模型很少生成針對底層硬件優(yōu)化的內(nèi)核，這表明未來模型仍有改進空間。

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

有趣的是，他們發(fā)現(xiàn)OpenAI o1和DeepSeek-R1生成的部分內(nèi)核使用了特定于硬件的指令和優(yōu)化。

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

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

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

結(jié)論

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

總的來說，盡管大多數(shù)基準(zhǔn)測試最終會達到飽和，但KernelBench被設(shè)計為隨著新的AI工作負載的出現(xiàn)而動態(tài)演進。

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

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

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

下一步工作

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

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

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

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

作者介紹

Anne Ouyang

Anne Ouyang目前是斯坦福大學(xué)計算機科學(xué)（CS）博士生，在Scaling Intelligence Lab（規(guī)?；悄軐嶒炇遥┻M行研究。

她的研究興趣主要集中在可擴展的自我改進機器學(xué)習(xí)系統(tǒng)，同時也廣泛關(guān)注實證機器學(xué)習(xí)（empirical ML）和性能工程（performance engineering）。

她曾獲得了MIT學(xué)士和碩士學(xué)位，并曾在NVIDIA cuDNN團隊工作，負責(zé)編寫CUDA內(nèi)核，用于加速GPU上的深度學(xué)習(xí)工作負載。

Simon Guo

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

他曾獲得了UC伯克利電氣工程和計算機科學(xué)學(xué)士學(xué)位。他的研究興趣在計算機系統(tǒng)和機器學(xué)習(xí)。

最近，他在Cohere從事語言模型預(yù)訓(xùn)練工作。在此之前，他曾在蘋果公司設(shè)計GPU，在Anyscale開發(fā)分布式系統(tǒng)，并在NVIDIA DRIVE部門從事自動駕駛汽車的開發(fā)工作。