DeepSeek R1不編程就能生成GPU內(nèi)核，比熟練工程師好，驚到了英偉達
DeepSeek 引爆 AI 社區(qū)后，人們都在嘗試本地部署和各領(lǐng)域應(yīng)用，在新模型基礎(chǔ)上持續(xù)改進的方向也被不斷提出。與此同時，英偉達卻在嘗試用 DeepSeek 給大模型 pipeline 本身搞自動化。
本周三，英偉達在博客中介紹了利用 DeepSeek-R1 和推理時擴展技術(shù)來自動生成優(yōu)化 GPU 內(nèi)核的最新研究成果，效果異常的好。
對此有人評價道：難不成英偉達在自拆護城河？
也有人已經(jīng)開始擔(dān)心自己的工作會不會被 AI 代替了。
隨著 AI 大模型規(guī)模不斷擴展，能力持續(xù)進步，測試時擴展（TTS：Test-Time Scaling）或推理時擴展（Inference-Time Scaling）法則正在興起。這項技術(shù)也被稱為 AI 推理或長思考，它通過在推理過程中分配額外的計算資源來評估多種可能的結(jié)果，然后選擇最佳的一個，從而提高模型整體性能。
推理能力的加強使得 AI 初步掌握了類似于人類剖析復(fù)雜問題的能力，能逐個解決以得出最終解決方案的方式，進行策略性思考和系統(tǒng)性地解決復(fù)雜問題。
在英偉達這篇文章中，工程師們進行了一項實驗，他們使用最新、最熱門的開源大模型 DeepSeek-R1 在推理過程中利用額外的計算能力來解決一個復(fù)雜問題 —— 自動生成數(shù)值正確，且針對不同注意力變體優(yōu)化的 GPU 注意力內(nèi)核，而無需任何顯式編程。
人們發(fā)現(xiàn)在某些情況下，R1 輸出的結(jié)果甚至優(yōu)于由熟練工程師開發(fā)出來的優(yōu)化內(nèi)核。
對優(yōu)化注意力內(nèi)核的需求及相關(guān)挑戰(zhàn)
注意力機制是徹底改變大型語言模型（LLM）發(fā)展的一個關(guān)鍵概念。它是一種強大的機制，使 AI 模型在執(zhí)行任務(wù)時能夠選擇性地關(guān)注輸入中最相關(guān)的部分。通過專注于重要信息，注意力操作幫助模型做出更好的預(yù)測并發(fā)現(xiàn)數(shù)據(jù)中的隱藏模式。
注意力操作的計算復(fù)雜度與輸入序列長度的平方成正比增長。這促使我們需要開發(fā)優(yōu)化的底層實現(xiàn)（即 GPU 內(nèi)核），以防止簡單實現(xiàn)導(dǎo)致的運行時錯誤（如內(nèi)存不足錯誤），并提高計算效率。
另外，注意力有多種變體（因果注意力、相對位置嵌入、ALiBi 等），工程師通常需要為特定任務(wù)組合使用這些變體。
多模態(tài)模型（例如視覺 Transformer）引入了額外的挑戰(zhàn)，因為它們需要專門的注意力機制（如空間鄰域注意力）來維護計算機視覺、視頻生成模型等中常見的時空信息。
圖 1：2D 輸入上的鄰域注意力。
但在這個任務(wù)上，即使對于經(jīng)驗豐富的軟件工程師來說，創(chuàng)建針對注意力的優(yōu)化 GPU 內(nèi)核也需要大量技能和時間。
最近的大模型（如 DeepSeek-R1）在代碼生成任務(wù)中表現(xiàn)出了很大的潛力，但它們在第一次嘗試創(chuàng)建優(yōu)化代碼時仍然效果不好。這使得在推理時使用其他策略來生成優(yōu)化代碼成為了必要。
以下 Prompt 是相對位置嵌入注意力內(nèi)核的示例用戶輸入。
復(fù)制
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.44269504
Please provide the complete updated kernel code that incorporates these changes, ensuring that the relative positional encoding is applied efficiently within the kernel operations.大模型有時會產(chǎn)生幻覺，或輸出混合不同語言或框架的語法，導(dǎo)致生成的代碼錯誤或效率低下。計算最佳 GPU 線程映射也是一項艱巨而具有挑戰(zhàn)性的任務(wù)，通常需要迭代細化才能獲得正確且高效的內(nèi)核。
用于生成優(yōu)化 GPU 內(nèi)核的推理時擴展
為了利用優(yōu)化的注意力內(nèi)核獲得最佳結(jié)果，英偉達工程師創(chuàng)建了一個新的工作流程，包括了一個特殊的驗證器以及一個在預(yù)定時間內(nèi)以閉環(huán)方式進行推理的 DeepSeek-R1 模型。
圖 2：在 Nvidia Hopper 平臺上利用 DeepSeek-R1 的推理時擴展。
具體地講，該工作流程首先由手動提示進行初始化，然后 DeepSeek-R1 在第一次遍歷中生成 GPU 代碼（即內(nèi)核）。驗證器在一塊英偉達 H100 GPU 上運行，它對生成的內(nèi)核進行分析，并創(chuàng)造新的提示以作為輸入提供給 DeepSeek-R1。
這種閉環(huán)方法每次都以不同的方式引導(dǎo)代碼生成過程，從而實現(xiàn)更好的效果。英偉達發(fā)現(xiàn)，這個過程持續(xù) 15 分鐘就可以得到一個改進的注意力內(nèi)核。
圖 3：使用 flex attention 自動生成的優(yōu)化注意力內(nèi)核的性能。
根據(jù)斯坦福的 KernelBench 基準(zhǔn)測試，該工作流程可以為 100% 的 Level-1 問題和 96% 的 Level-2 問題生成數(shù)值正確的內(nèi)核。
其中，KernelBench 中的 Level-1 解決率是指：評估 LLM 為特定計算任務(wù)生成高效 GPU 內(nèi)核的能力的數(shù)值正確指標(biāo)。該測試是一系列「測試最新 LLM GPU 編程能力」挑戰(zhàn)的一部分。
圖 4 顯示了推理時預(yù)算對 agent 解決率的影響，結(jié)果顯示，在 Level-1 類別中為每個問題分配超過 10 分鐘的時間，可以使工作流程為 100 個問題中的大多數(shù)生成數(shù)值正確的代碼。
這些結(jié)果表明，在使用最新的 DeepSeek-R1 模型時，如果在推理階段投入更多計算，則可以獲得更好的 GPU 內(nèi)核。英偉達表示，其對于 DeepSeek-R1 的最新進展及其應(yīng)用潛力感到非常興奮。
本文轉(zhuǎn)自機器之心 ，作者：機器之心
