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 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.

大模型有時會產生幻覺，或輸出混合不同語言或框架的語法，導致生成的代碼錯誤或效率低下。計算最佳 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 的最新進展及其應用潛力感到非常興奮。