隨著 AI 模型擴展其功能以解決更復雜的挑戰，一種稱為“ 測試時擴展 ”或“ 推理時擴展 ”的新擴展法則正在出現。該技術也稱為 AI 推理 或 長時思考 技術，通過在推理過程中分配額外的計算資源來評估多種可能的結果，然后選擇最佳結果 (neural network)，從而提高模型性能。這使得 AI 能夠以類似于人類剖析復雜問題并單獨解決這些問題以達成最終解決方案的方式，制定戰略并系統化地解決復雜問題。
在本文中，我們將介紹 NVIDIA 工程師完成的一項實驗，他們在推理過程中使用最新的開源模型之一 DeepSeek-R1 模型以及額外的計算能力來解決復雜的問題。該實驗旨在自動生成 GPU 注意力內核，這些內核在數值上是正確的，并針對不同的注意力類型進行了優化，而無需任何顯式編程。
事實證明，在某些情況下，最終結果優于由技術精湛的工程師開發的優化內核。
優化注意力內核的需求和相關挑戰
注意力是一個關鍵概念，徹底改變了大語言模型（LLM）的發展。這是一種功能強大的機制，可讓 AI 模型在執行任務時，有選擇地專注于與輸入內容最相關的部分。通過聚焦于重要信息，注意力運算可幫助模型做出更好的預測，并找到數據中隱藏的模式。
相對于輸入序列長度，注意力運算的計算復雜性呈正交增長。這促使我們需要開發經過優化的低級實現 (即 GPU 內核)，以防止簡單實現產生的運行時錯誤 (例如內存不足的錯誤)，并提高計算效率。
注意力有多種變體 (因果關系、相對位置嵌入、不在場證明等)，工程師通常必須將這些變體的組合用于給定任務。
多模態模型 (例如，vision transformers) 帶來了額外的一層挑戰，因為它們需要專門的注意力機制 (Spatial Neighborhood Attention) 來維護計算機視覺、視頻生成模型等領域中經常遇到的時空信息。
創建經優化的 GPU 內核以供關注需要大量技能和時間，即使對于經驗豐富的軟件工程師而言也是如此。
最近的 LLMs（如 DeepSeek-R1）在代碼生成任務方面表現出了很大的前景，但在第一次嘗試創建優化代碼時，它們仍然面臨挑戰。這使得有必要在推理時使用其他策略來生成優化的代碼。
以下提示是用戶輸入相對位置嵌入注意力核函數的示例。
Please write a GPU attention kernel to support relative position encodings. Implement the relative positional encoding on the fly within the kernel. The complete code should be returned, including the necessary modifications.Use the following function to compute the relative positional encoding:def relative_positional(score, b, h, q_idx, kv_idx): return score + (q_idx - kv_idx)When implementing the kernel, keep in mind that a constant scaling factor 1.44269504 should be applied to the relative positional encoding due to qk_scale = sm_scale * 1.44269504. The PyTorch reference does not need to scale the relative positional encoding, but in the GPU kernel, use:qk = qk * qk_scale + rel_pos * 1.44269504Please provide the complete updated kernel code that incorporates these changes, ensuring that the relative positional encoding is applied efficiently within the kernel operations. |
LLM 偶爾會產生來自不同語言或框架的幻影代碼或混合語法，從而立即導致代碼錯誤或效率低下。計算最佳 GPU 線程映射也并非易事，也是一項具有挑戰性的任務，通常需要進行迭代優化才能實現正確高效的內核。
用于生成經過優化的 GPU 內核的推理時間擴展
為了通過優化的注意力內核獲得最佳結果，NVIDIA 工程師創建了一個新的工作流程，其中包括一個特殊的驗證器以及 DeepSeek-R1 模型，在預先設定的時間內以閉環方式進行推理。
該工作流程首先通過手動提示進行初始化，然后 DeepSeek-R1 模型會在首次通道中生成 GPU 代碼(即核函數)。該驗證器在 NVIDIA H100 GPU 上運行。它會分析生成的核函數并創建新的提示，這些提示以 input 的形式提供給 DeepSeek-R1 模型。
這種閉環方法每次都以不同的方式指導代碼生成過程，從而改進代碼生成過程。該團隊發現，讓此過程持續 15 分鐘可以改進注意力核函數。
此工作流程為 100%的 1 級問題和 96%的 2 級問題生成了數值正確的內核，測試對象為 斯坦福大學的 KernelBench 基準測試。*
KernelBench 中的 1 級求解率是指用于評估 LLM 為特定計算任務生成高效 GPU 內核的能力的數字正確指標。本次測試屬于一系列挑戰，旨在測試最新 LLM 在 GPU 編程中的能力。
圖 4 顯示了推理時間預算如何影響智能體的求解率。在 Level-1 類別中為每個問題分配超過 10 分鐘的時間，使工作流程能夠為 100 個問題中的大多數生成正確的數字代碼。
DeepSeek-R1 上經過優化的 GPU 內核
這些結果展示了如何使用最新的 DeepSeek-R1 模型，通過在推理期間使用更強大的計算能力來提供更好的 GPU 內核。這仍然是一個新的研究領域，在自動生成有效注意力內核的前景良好的方法方面取得了早期成果。
雖然我們已經有了一個良好的開端，但我們需要做更多的工作，以便為更廣泛的問題持續提供更好的結果。我們對 DeepSeek-R1 的最新進展及其潛力感到興奮。
如需了解更多信息或入門，請參閱 DeepSeek-R1 NIM 微服務 (現已在 build.nvidia.com 上提供)。
?
?
