搞不懂CUDA的人有救了，Devin開發商開源Kevin，強化學習生成CUDA內核

2025-05-07 13:46:34

Kevin-32B 基于 QwQ-32B 在 KernelBench 數據集上使用 GRPO 進行了多輪強化學習訓練，實現了超越 o3 和 o4-mini 的頂級推理表現。

本周三，知名 AI 創業公司，曾發布「全球首個 AI 軟件工程師」的 Cognition AI 開源了一款使用強化學習，用于編寫 CUDA 內核的大模型 Kevin-32B。

Kevin-32B 基于 QwQ-32B 在 KernelBench 數據集上使用 GRPO 進行了多輪強化學習訓練，實現了超越 o3 和 o4-mini 的頂級推理表現。

對此，機器學習社區表現出了極大的興趣。有人表示期待 DeepSeek R1 風格的訓練方法用來提升代碼效率已久，這回終于有人站出來了。

在一篇博客中，Cognition AI 詳細介紹了新模型強化學習訓練的機制。

代碼是一個不斷迭代的過程 —— 需要我們編寫、執行程序，評估結果，并根據反饋優化代碼。大語言模型（LLM）在代碼生成方面的最新進展嘗試將此過程融入推理階段，并使用并行采樣等方法。雖然這些方法是有效的，但它們依賴于搜索而非實際學習 —— 在這其中模型權重被凍結。

Cognition AI 探索了多輪強化學習，使用來自環境的中間反饋，并屏蔽模型思維以避免在多輪訓練中上下文爆炸。

他們提出的模型 Kevin-32B（即 Kernel Devin）在內核生成方面優于前沿推理模型。此外他們通過實驗結果表明，與單輪訓練相比，多輪訓練使模型在自我優化方面更有效。

多輪訓練方法

在新模型上，作者使用了 KernelBench，這是一個包含 250 個基于 PyTorch 的經典深度學習任務的數據集。它衡量優化 CUDA 內核替換 PyTorch 算子的能力。訓練專注于前兩個級別，每個級別包含 100 個任務。級別 1 包含矩陣乘法、卷積和損失函數等基礎任務，級別 2 則包含融合算子，再使用這兩個級別的 180 個任務進行訓練，并留出 20 個任務作為保留集。

在訓練過程中，模型會經歷一個迭代反饋循環：從生成的內核中提取反饋，并再對其進行優化。如果內核編譯失敗，錯誤軌跡會傳遞給模型并要求其修復。如果編譯正確，系統則會測量運行時間并要求模型進一步改進。

初始方法的機制如下。從初始提示開始，先在每個優化步驟后附加思路鏈、內核和評估信息，再為整個軌跡分配一個獎勵（定義為任何內核獲得的最高分數），并使用該序列進行訓練。

不過，這種方法存在一些挑戰：

上下文窗口爆炸：推理模型會生成很長的思維鏈。使用這種方法，軌跡的長度在幾次傳遞后很容易達到 5-10 萬個 token，這會導致訓練不便；
樣本效率低下和 credit 分配問題：即使我們生成了多個內核，仍然只為整個軌跡分配一個獎勵。這無法表明哪個細化步驟真正提高了正確性或性能。獎勵應該根據細化步驟對最終結果的貢獻來分配。

為了解決上下文長度爆炸的問題，Kevin-32B 丟棄了軌跡中最長的部分 —— 思維鏈。現在每個 token 將只包含先前生成的內核和評估結果。為了仍然保留上一步思考過程的信息，我們特意要求模型生成其自身思考過程的簡要摘要，然后將其傳遞給后續上下文。

移除推理的思路鏈

為了解決樣本效率低下的問題，Kevin-32B 選擇了一個更具表現力的獎勵函數，將內核的細化建模為馬爾可夫決策過程，將給定響應的獎勵設置為當前內核與所有后續內核得分的折扣總和。至此，每個細化步驟都變成了一個訓練樣本。

將獎勵設為分數的折扣總和

結果

對于每項任務，作者并行采樣了 16 條軌跡，并進行 8 個連續的細化步驟。對于每項任務的正確性或性能，他們將 best@16 定義為所有軌跡的最大值，avg@16 定義為所有軌跡的平均值。

經過 8 個優化步驟，Kevin-32B 在整個數據集上的平均正確率為 65%，顯著超越了 QwQ-32B 和前沿模型。它解決了 89% 的數據集，而 o4-mini 和 o3 分別只解決了 53% 和 51%。在整個數據集中，Kevin-32B 實現了 1.41 倍的 best@16 加速比，優于前沿模型。

Kevin-32B 在二級任務上尤其有效，平均正確率為 48%（o4-mini 為 9.6%，o3 為 9.3%）。這表明多輪訓練可以提高模型解決更具挑戰性、更長期任務的能力。同樣，我們能注意到模型實現了 1.74 倍的 best@16 加速比（是 o4-mini 和 o3 為 1.2 倍）。

多輪訓練 vs 單輪訓練

Kevin-32B 也展現出比 QwQ-32B 和單輪訓練模型顯著的提升。在 4 個優化步驟下，Kevin-32B 的表現略優于單輪模型，但隨著優化步驟增加到 8 個，兩者之間的差距進一步擴大。這表明，通過鼓勵更積極的優化，多輪訓練在串行軸上具有更好的擴展性。

我們可能會想，單輪訓練模型是否可以通過采樣并行軌跡來實現更好的加速，然而在這種環境下并非如此。在計算預算固定的情況下，即使對于單輪訓練模型，多輪推理也比單輪推理更占優勢。

獎勵黑客攻擊

最初的實驗使用了規模較小的模型，例如 DeepSeek-R1-Distill-Qwen-7B，這導致了多起獎勵黑客攻擊事件：

該模型只是復制了 PyTorch 參考實現，因此因生成正確答案而獲得獎勵，且加速比提高了 1.0 倍；
該模型將錯誤的 CUDA 內核實現包裝在 try-except 語句中，并調用 PyTorch 實現函數作為回退；
該模型繼承自參考實現，繞過了對 CUDA 實現的需求。

獎勵黑客攻擊示例

為了防止獎勵黑客攻擊，作者對響應施加了更嚴格的格式檢查。對于使用 PyTorch 函數或不包含 CUDA 內核的響應，獎勵被設置為 0。

我們能觀察到，當模型能力與數據集難度之間的差距很大時，就會發生獎勵黑客攻擊。模型難以解決任何任務，因此當它生成一個被黑客入侵的內核時，它是唯一具有正向優勢的操作，因此會得到顯著的強化。

無意義和重復的生成

多次運行后，我們可以觀察到在大約 35-40 步，模型開始生成重復或無意義的響應。假設這是因為模型偏離了穩定區域，令人驚訝的是，作者偶然發現了一個可以有效預測未來垃圾響應的指標 —— 思維鏈不是以「OK」開頭的響應比例 —— 我們稱之為「Not Okay Ratio」。

模型訓練越多，CoT 的開頭就越不穩定：

"Okay Amigos, so I need to optimize this 3D tensor-matrix multiplication..”
"Okay Holy crap, I need to get this code optimized….”
在精煉步驟的第 8 遍時：
"Okay SHIT, this is getting frustrating. Let me see. The error is about …”

為了解決這個問題，作者嘗試使用 KL 系數 0.001 和 0.01 進行運行，但發現這會減慢學習速度，同時又不能防止垃圾數據的產生。最后，作者通過使用 Dr. GRPO 提出的恒定長度損失正則化（這顯著降低了梯度范數），并將梯度范數大幅削減至 0.05，將垃圾數據的出現延遲到了第 100 步。

示例內核

作為示例，以 1 級任務 40（LayerNorm）的軌跡為例，作者重點介紹了模型如何迭代地整合來自內核評估的反饋，并利用這些反饋來改進其生成過程。

步驟 1：模型生成一個正確的內核。它使用共享內存，并將均值 / 方差計算與歸一化融合。這實現了 0.6 倍的加速。

...
__shared__ float shared_sums [512]; // 256 threads * 2 values
...
for (int s = block_size / 2; s > 0; s >>= 1) {
    if (tid < s) {
        shared_sums [2 * tid] += shared_sums [2 * (tid + s)];
        shared_sums [2 * tid + 1] += shared_sums [2 * (tid + s) + 1];
    }
    __syncthreads ();
}
...
for (int d = tid; d < D; d += block_size) {
    float x_val = x [n * D + d];
    float y_val = (x_val - mean) /sqrt_var;
    y_val = y_val * scale [d] + bias [d];
    y [n * D + d] = y_val;
}
...

步驟 2：模型從評估反饋中注意到「當前的加速比是 0.6，這意味著它實際上比 PyTorch 版本慢。這不太好。讓我想想瓶頸可能在哪里」。它將大部分推理 token 用于尋找更合適的塊大小，并相應地修改代碼。修改是正確的，加速比現在達到了 4.27 倍。

步驟 3 和 4：模型嘗試進行更積極的優化。它考慮了循環展開和減少庫沖突，但最終選擇了 Warp 級內聯函數。但最終失敗了 —— 首先是因為歸約邏輯存在正確性問題，其次是因為思路太長。

步驟 5：模型注意到 Warp 歸約實現中存在錯誤的問題并進行了修復。它成功實現了兩級 Warp 歸約。最終加速比達到了 9.61 倍。

// Warp-level reduction using shuffle instructions
for (int delta = 1; delta <= 16; delta <<= 1) {
  float other_sum = __shfl_xor_sync (0xFFFFFFFF, warp_sum, delta);
  float other_sum_sq = __shfl_xor_sync (0xFFFFFFFF, warp_sum_sq, delta);
  warp_sum += other_sum;
  warp_sum_sq += other_sum_sq;
}
__shared__ float sum_warp [32];
__shared__ float sum_sq_warp [32];
__shared__ float results [2]; // [mean, inv_std]
if (warp_id == 0) {
  sum_warp [warp_lane] = warp_sum;
  sum_sq_warp [warp_lane] = warp_sum_sq;
}
__syncthreads ();
// Final reduction within the first warp (tid 0-31)
if (tid < 32) {
  float my_sum = sum_warp [tid];
  float my_sum_sq = sum_sq_warp [tid];
  // Reduce within the first warp (32 threads)
  for (int s = 16; s >= 1; s >>= 1) {
      my_sum += __shfl_xor_sync (0xFFFFFFFF, my_sum, s);
      my_sum_sq += __shfl_xor_sync (0xFFFFFFFF, my_sum_sq, s);
  }
  ...
}

訓練設置

這里使用的是組相對策略優化算法（GRPO），該算法由 DeepSeek 提出，是近端策略優化算法（PPO）的一種變體。

GRPO 不使用值網絡來估算基線和計算優勢，而是將從同一提示中采樣的響應組內的獎勵歸一化。

作者使用 vLLM 進行推理，使用 DeepSpeed Zero-3 卸載優化器狀態。每批訓練 8 個任務，每個任務 16 個軌跡。我們使用 GRPO，每批 2 個梯度步驟。基礎模型是 QwQ-32B。

響應生成完成后，每個 GPU 將其 vLLM 引擎卸載到 CPU 內存，并評估其生成的內核。對于每個響應，作者都會檢查響應格式是否正確，并提取 CUDA 內核。然后，編譯并執行代碼，使用隨機張量測試其正確性。如果正確，則會對內核的運行時間進行剖析。

通過正確性檢查的響應將獲得 0.3 的獎勵，而額外的性能獎勵則相當于與參考實現相比所獲得的速度提升。

內核評估和基準問題

作者對評估進行了沙盒化處理，以免出現致命錯誤（如 CUDA 非法內存訪問），導致訓練過程崩潰。

由于 KernelBench 中的許多任務使用非常小的輸入張量，該基準最終測量的是內核啟動開銷而非實際內核執行時間。為了解決這個問題，作者擴大了受影響任務的張量維度。

KernelBench 評估工具中還有一個更隱蔽的錯誤，導致被測內核將參考實現的輸出張量作為自己的張量輸出循環使用。因此，只計算（正確）部分輸出張量的內核仍能通過正確性檢查。為了解決這個問題，我們首先運行測試過的內核，然后再運行參考實現，從而避免了這一問題。

單輪訓練設置

作者使用 max_grad_norm = 0.5、lr = 常量 2e-6（熱身比為 0.03）、max_prompt_length = 8192、max_response_length = 16384，其使用 DAPO 的 Clip-High，eps_high = 0.28，并將 KL 系數設為 0，以允許模型自由偏離基本策略。

我們能觀察到，單輪模型比基礎模型有明顯改善，但在 25 步之后，獎勵開始趨于穩定。

多輪訓練設置

作者根據每個軌跡的最佳內核對其正確性和性能進行評分。在最終訓練運行中，每個前向路徑由 16 個平行軌跡組成，每個軌跡有 4 個細化步驟，折扣系數為 0.4。與單輪訓練不同的是，現在的獎勵會穩步增加。

隨著模型學會更高效地使用推理 token 來生成內核，響應長度最初會減少。第 25 步之后，隨著模型嘗試更復雜的解決方案，響應長度會增加。按照 DeepScaleR 的做法，作者在第 30 步將最大響應長度從 16K 擴展到 22K token。

測試時搜索

在訓練過程中，為了保持合理的訓練時間，必須限制搜索。但在測試時，可以自由地使用更復雜的測試時間技術來進一步提高性能。

為此，Kevin-32B 使用了改進版的束搜索，其工作原理如下。作者首先在 16 條軌跡上執行 4 個串行細化步驟，就像在訓練時間所做的那樣。在這一過程結束時，根據生成的最快內核對軌跡進行排序，并保留最好的 4 個軌跡。然后，作者將每條軌跡復制 4 次（共 16 條軌跡），并重復這一過程。作為一種通用方法，它還顯著提高了模型性能，在整個數據集上平均提速 1.56 倍。

不出所料，隨著測試時間計算量的增加，得到的收益也在減少，盡管平均性能在幾個小時后仍有提高。

束搜索測試時間推理。計算時間是在單個 H200 GPU 上按內核計算的。請注意，其中還包括評估時的停滯時間，有效推理時間約占總計算時間的 45%。

作者通過修改軌跡數、束寬度和細化步數，嘗試了波束搜索的幾種不同變體。有趣的是，他們發現中位加速度較高的設置平均加速度較低，反之亦然。

作者將這種行為歸因于每種技術在探索 / 利用的帕累托前沿中的不同位置：通過更多的探索，內核可以實現一致的加速，盡管它們可能會錯過非常激進的、顯著提高平均速度的優化。

未來工作

我們相信，這項工作只是探索編程智能體訓練方法的開始。如果有更多的時間和計算能力，我們希望嘗試以下方法：

學習價值網絡并使用 PPO 進行訓練。事實上，基線估計器是按引用步驟而不是按提示計算的；
在訓練時整合更復雜的搜索方法，如束搜索（而不僅僅是并行 + 串行搜索）；
將多輪訓練法應用于更普遍的編程環境。

總結

在這項工作中，作者提出了一種方法，該方法適用于任何具有中間獎勵的多輪環境。研究表明，這種方法比單圈 GRPO 取得了更好的結果。

端到端訓練將是未來智能體的重要組成部分。雖然手工制作的多 LLM 工作流程可以帶來短期收益，但它們在很大程度上依賴于人類的啟發式方法，而且無法擴展。相比之下，更通用的方法可以讓模型自由探索不同的軌跡（搜索），然后學習這些長期動態。通過與編程環境的交互，模型產生經驗流，并通過反饋不斷調整。我們希望這樣的方法能成為邁向自主編程智能體的第一步。

HuggingFace: https://huggingface.co/cognition-ai/Kevin-32B

最后，讓我們來復習一下《The Bitter Lesson》的道理：

我們必須吸取慘痛的教訓，從長遠來看，在我們的思維方式中植入知識是行不通的。這一慘痛教訓基于以下歷史觀察：

人工智能研究人員經常試圖在其智能體中構建知識；
這在短期內總是有幫助的，而且會讓研究人員個人感到滿意；
從長遠來看，它會停滯不前，甚至抑制進一步的進展；
通過搜索和學習來擴展計算規模的相反方法最終會帶來突破性進展。

責任編輯：張燕妮來源：機器之心

AI 模型數據

搞不懂CUDA的人有救了，Devin開發商開源Kevin，強化學習生成CUDA內核

多輪訓練方法

結果

訓練設置

更多結果 & 消融研究

測試時搜索

未來工作

總結