
社群是國內外知名的機器學習與自然語言處理社群,受眾覆蓋國內外NLP碩博生、高校老師以及企業研究人員。
來源 | 機器之心
編輯|澤南、杜偉
這項嘗試只用到了 R1 模型和基本驗證器,沒有針對 R1 的工具,沒有對專有的英偉達程式碼進行微調。其實根據 DeepSeek 介紹,R1 的編碼能力不算頂尖。
DeepSeek 引爆 AI 社群後,人們都在嘗試本地部署和各領域應用,在新模型基礎上持續改進的方向也被不斷提出。與此同時,英偉達卻在嘗試用 DeepSeek 給大模型 pipeline 本身搞自動化。
本週三,英偉達在部落格中介紹了利用 DeepSeek-R1 和推理時擴充套件技術來自動生成最佳化 GPU 核心的最新研究成果,效果異常的好。

對此有人評價道:難不成英偉達在自拆護城河?

也有人已經開始擔心自己的工作會不會被 AI 代替了。

隨著 AI 大模型規模不斷擴充套件,能力持續進步,測試時擴充套件(TTS:Test-Time Scaling)或推理時擴充套件(Inference-Time Scaling)法則正在興起。這項技術也被稱為 AI 推理或長思考,它透過在推理過程中分配額外的計算資源來評估多種可能的結果,然後選擇最佳的一個,從而提高模型整體效能。
推理能力的加強使得 AI 初步掌握了類似於人類剖析複雜問題的能力,能逐個解決以得出最終解決方案的方式,進行策略性思考和系統性地解決複雜問題。
在英偉達這篇文章中,工程師們進行了一項實驗,他們使用最新、最熱門的開源大模型 DeepSeek-R1 在推理過程中利用額外的計算能力來解決一個複雜問題 —— 自動生成數值正確,且針對不同注意力變體最佳化的 GPU 注意力核心,而無需任何顯式程式設計。
人們發現在某些情況下,R1 輸出的結果甚至優於由熟練工程師開發出來的最佳化核心。
對最佳化注意力核心的需求及相關挑戰
注意力機制是徹底改變大型語言模型(LLM)發展的一個關鍵概念。它是一種強大的機制,使 AI 模型在執行任務時能夠選擇性地關注輸入中最相關的部分。透過專注於重要資訊,注意力操作幫助模型做出更好的預測並發現資料中的隱藏模式。
注意力操作的計算複雜度與輸入序列長度的平方成正比增長。這促使我們需要開發最佳化的底層實現(即 GPU 核心),以防止簡單實現導致的執行時錯誤(如記憶體不足錯誤),並提高計算效率。
另外,注意力有多種變體(因果注意力、相對位置嵌入、ALiBi 等),工程師通常需要為特定任務組合使用這些變體。
多模態模型(例如視覺 Transformer)引入了額外的挑戰,因為它們需要專門的注意力機制(如空間鄰域注意力)來維護計算機視覺、影片生成模型等中常見的時空資訊。

圖 1:2D 輸入上的鄰域注意力。
但在這個任務上,即使對於經驗豐富的軟體工程師來說,建立針對注意力的最佳化 GPU 核心也需要大量技能和時間。
最近的大模型(如 DeepSeek-R1)在程式碼生成任務中表現出了很大的潛力,但它們在第一次嘗試建立最佳化程式碼時仍然效果不好。這使得在推理時使用其他策略來生成最佳化程式碼成為了必要。
以下 Prompt 是相對位置嵌入注意力核心的示例使用者輸入。
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 followingfunctiontocompute the relative positional encoding:
def relative_positional(score, b, h, q_idx, kv_idx):
return score + (q_idx - kv_idx)
When implementing the kernel, keepin 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.44269504
Please provide the completeupdated kernel code that incorporates these changes, ensuring that the relative positional encodingis applied efficiently within the kernel operations.
大模型有時會產生幻覺,或輸出混合不同語言或框架的語法,導致生成的程式碼錯誤或效率低下。計算最佳 GPU 執行緒對映也是一項艱鉅而具有挑戰性的任務,通常需要迭代細化才能獲得正確且高效的核心。
用於生成最佳化 GPU 核心的推理時擴充套件
為了利用最佳化的注意力核心獲得最佳結果,英偉達工程師建立了一個新的工作流程,包括了一個特殊的驗證器以及一個在預定時間內以閉環方式進行推理的 DeepSeek-R1 模型。

圖 2:在 Nvidia Hopper 平臺上利用 DeepSeek-R1 的推理時擴充套件。
具體地講,該工作流程首先由手動提示進行初始化,然後 DeepSeek-R1 在第一次遍歷中生成 GPU 程式碼(即核心)。驗證器在一塊英偉達 H100 GPU 上執行,它對生成的核心進行分析,並創造新的提示以作為輸入提供給 DeepSeek-R1。
這種閉環方法每次都以不同的方式引導程式碼生成過程,從而實現更好的效果。英偉達發現,這個過程持續 15 分鐘就可以得到一個改進的注意力核心。

圖 3:使用 flex attention 自動生成的最佳化注意力核心的效能。
根據斯坦福的 KernelBench 基準測試,該工作流程可以為 100% 的 Level-1 問題和 96% 的 Level-2 問題生成數值正確的核心。
其中,KernelBench 中的 Level-1 解決率是指:評估 LLM 為特定計算任務生成高效 GPU 核心的能力的數值正確指標。該測試是一系列「測試最新 LLM GPU 程式設計能力」挑戰的一部分。
圖 4 顯示了推理時預算對 agent 解決率的影響,結果顯示,在 Level-1 類別中為每個問題分配超過 10 分鐘的時間,可以使工作流程為 100 個問題中的大多數生成數值正確的程式碼。

這些結果表明,在使用最新的 DeepSeek-R1 模型時,如果在推理階段投入更多計算,則可以獲得更好的 GPU 核心。英偉達表示,其對於 DeepSeek-R1 的最新進展及其應用潛力感到非常興奮。
參考內容:https://developer.nvidia.com/blog/automating-gpu-kernel-generation-with-deepseek-r1-and-inference-time-scaling/
技術交流群邀請函

△長按新增小助手
掃描二維碼新增小助手微信
關於我們
