新智元報道
新智元報道
編輯：編輯部
【新智元導讀】斯坦福和普林斯頓研究者發現，DeepSeek-R1生成的自定義CUDA核心，完爆了o1和Claude 3.5 Sonnet，拿下總排名第一。雖然目前只能在不到20%任務上超越PyTorch Eager基線，但GPU程式設計加速自動化的按鈕，已經被按下！
近日，來自斯坦福和普林斯頓的研究者發現，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基線。
這些模型生成的核心存在大量執行錯誤、功能正確性問題，並且無法進行特定平臺的最佳化。
具體來說，研究者發現：
-
對模型而言，編寫功能正確的核心仍然具有挑戰性； -
模型透過最佳化展示了生成高效能核心的潛力； -
利用反饋對於減少執行錯誤和發現更快的方案很重要。
當然，KernelBench目前還只是讓GPU加速奔跑的起點，但也是讓整個GPU程式設計自動化的起始催化劑。
令人興奮的是，現在已經有了許多新的工作，專注於解決KernelBench中涉及的問題。
隨後在2月12日，Meta也發文測試了前沿模型編寫GPU核心方面的效能，他們發現，最佳模型可以在KernelBench上提供平均1.8倍的加速。
如雨後春筍般出現的研究表明，如今，我們已經進入了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部門從事自動駕駛汽車的開發工作。
參考資料：HNYZs
https://arxiv.org/abs/2502.10517
https://x.com/simonguozirui/status/1894455889039692263