DeepSeek 引爆 AI 社區(qū)后，人們都在嘗試本地部署和各領(lǐng)域應(yīng)用，在新模型基礎(chǔ)上持續(xù)改進(jìn)的方向也被不斷提出。與此同時(shí)，英偉達(dá)卻在嘗試用 DeepSeek 給大模型 pipeline 本身搞自動(dòng)化。

本周三，英偉達(dá)在博客中介紹了利用 DeepSeek-R1 和推理時(shí)擴(kuò)展技術(shù)來(lái)自動(dòng)生成優(yōu)化 GPU 內(nèi)核的最新研究成果，效果異常的好。

對(duì)此有人評(píng)價(jià)道：難不成英偉達(dá)在自拆護(hù)城河？

也有人已經(jīng)開(kāi)始擔(dān)心自己的工作會(huì)不會(huì)被 AI 代替了。

隨著 AI 大模型規(guī)模不斷擴(kuò)展，能力持續(xù)進(jìn)步，測(cè)試時(shí)擴(kuò)展（TTS：Test-Time Scaling）或推理時(shí)擴(kuò)展（Inference-Time Scaling）法則正在興起。這項(xiàng)技術(shù)也被稱為 AI 推理或長(zhǎng)思考，它通過(guò)在推理過(guò)程中分配額外的計(jì)算資源來(lái)評(píng)估多種可能的結(jié)果，然后選擇最佳的一個(gè)，從而提高模型整體性能。

推理能力的加強(qiáng)使得 AI 初步掌握了類似于人類剖析復(fù)雜問(wèn)題的能力，能逐個(gè)解決以得出最終解決方案的方式，進(jìn)行策略性思考和系統(tǒng)性地解決復(fù)雜問(wèn)題。

在英偉達(dá)這篇文章中，工程師們進(jìn)行了一項(xiàng)實(shí)驗(yàn)，他們使用最新、最熱門(mén)的開(kāi)源大模型 DeepSeek-R1 在推理過(guò)程中利用額外的計(jì)算能力來(lái)解決一個(gè)復(fù)雜問(wèn)題 —— 自動(dòng)生成數(shù)值正確，且針對(duì)不同注意力變體優(yōu)化的 GPU 注意力內(nèi)核，而無(wú)需任何顯式編程。

人們發(fā)現(xiàn)在某些情況下，R1 輸出的結(jié)果甚至優(yōu)于由熟練工程師開(kāi)發(fā)出來(lái)的優(yōu)化內(nèi)核。

對(duì)優(yōu)化注意力內(nèi)核的需求及相關(guān)挑戰(zhàn)

注意力機(jī)制是徹底改變大型語(yǔ)言模型（LLM）發(fā)展的一個(gè)關(guān)鍵概念。它是一種強(qiáng)大的機(jī)制，使 AI 模型在執(zhí)行任務(wù)時(shí)能夠選擇性地關(guān)注輸入中最相關(guān)的部分。通過(guò)專注于重要信息，注意力操作幫助模型做出更好的預(yù)測(cè)并發(fā)現(xiàn)數(shù)據(jù)中的隱藏模式。

注意力操作的計(jì)算復(fù)雜度與輸入序列長(zhǎng)度的平方成正比增長(zhǎng)。這促使我們需要開(kāi)發(fā)優(yōu)化的底層實(shí)現(xiàn)（即 GPU 內(nèi)核），以防止簡(jiǎn)單實(shí)現(xiàn)導(dǎo)致的運(yùn)行時(shí)錯(cuò)誤（如內(nèi)存不足錯(cuò)誤），并提高計(jì)算效率。

另外，注意力有多種變體（因果注意力、相對(duì)位置嵌入、ALiBi 等），工程師通常需要為特定任務(wù)組合使用這些變體。

多模態(tài)模型（例如視覺(jué) Transformer）引入了額外的挑戰(zhàn)，因?yàn)樗鼈冃枰獙ｉT(mén)的注意力機(jī)制（如空間鄰域注意力）來(lái)維護(hù)計(jì)算機(jī)視覺(jué)、視頻生成模型等中常見(jiàn)的時(shí)空信息。

圖 1：2D 輸入上的鄰域注意力。

但在這個(gè)任務(wù)上，即使對(duì)于經(jīng)驗(yàn)豐富的軟件工程師來(lái)說(shuō)，創(chuàng)建針對(duì)注意力的優(yōu)化 GPU 內(nèi)核也需要大量技能和時(shí)間。

最近的大模型（如 DeepSeek-R1）在代碼生成任務(wù)中表現(xiàn)出了很大的潛力，但它們?cè)诘谝淮螄L試創(chuàng)建優(yōu)化代碼時(shí)仍然效果不好。這使得在推理時(shí)使用其他策略來(lái)生成優(yōu)化代碼成為了必要。

以下 Prompt 是相對(duì)位置嵌入注意力內(nèi)核的示例用戶輸入。

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.

大模型有時(shí)會(huì)產(chǎn)生幻覺(jué)，或輸出混合不同語(yǔ)言或框架的語(yǔ)法，導(dǎo)致生成的代碼錯(cuò)誤或效率低下。計(jì)算最佳 GPU 線程映射也是一項(xiàng)艱巨而具有挑戰(zhàn)性的任務(wù)，通常需要迭代細(xì)化才能獲得正確且高效的內(nèi)核。

用于生成優(yōu)化 GPU 內(nèi)核的推理時(shí)擴(kuò)展

為了利用優(yōu)化的注意力內(nèi)核獲得最佳結(jié)果，英偉達(dá)工程師創(chuàng)建了一個(gè)新的工作流程，包括了一個(gè)特殊的驗(yàn)證器以及一個(gè)在預(yù)定時(shí)間內(nèi)以閉環(huán)方式進(jìn)行推理的 DeepSeek-R1 模型。

圖 2：在 Nvidia Hopper 平臺(tái)上利用 DeepSeek-R1 的推理時(shí)擴(kuò)展。

具體地講，該工作流程首先由手動(dòng)提示進(jìn)行初始化，然后 DeepSeek-R1 在第一次遍歷中生成 GPU 代碼（即內(nèi)核）。驗(yàn)證器在一塊英偉達(dá) H100 GPU 上運(yùn)行，它對(duì)生成的內(nèi)核進(jìn)行分析，并創(chuàng)造新的提示以作為輸入提供給 DeepSeek-R1。

這種閉環(huán)方法每次都以不同的方式引導(dǎo)代碼生成過(guò)程，從而實(shí)現(xiàn)更好的效果。英偉達(dá)發(fā)現(xiàn)，這個(gè)過(guò)程持續(xù) 15 分鐘就可以得到一個(gè)改進(jìn)的注意力內(nèi)核。

圖 3：使用 flex attention 自動(dòng)生成的優(yōu)化注意力內(nèi)核的性能。

根據(jù)斯坦福的 KernelBench 基準(zhǔn)測(cè)試，該工作流程可以為 100% 的 Level-1 問(wèn)題和 96% 的 Level-2 問(wèn)題生成數(shù)值正確的內(nèi)核。

其中，KernelBench 中的 Level-1 解決率是指：評(píng)估 LLM 為特定計(jì)算任務(wù)生成高效 GPU 內(nèi)核的能力的數(shù)值正確指標(biāo)。該測(cè)試是一系列「測(cè)試最新 LLM GPU 編程能力」挑戰(zhàn)的一部分。

圖 4 顯示了推理時(shí)預(yù)算對(duì) agent 解決率的影響，結(jié)果顯示，在 Level-1 類別中為每個(gè)問(wèn)題分配超過(guò) 10 分鐘的時(shí)間，可以使工作流程為 100 個(gè)問(wèn)題中的大多數(shù)生成數(shù)值正確的代碼。

這些結(jié)果表明，在使用最新的 DeepSeek-R1 模型時(shí)，如果在推理階段投入更多計(jì)算，則可以獲得更好的 GPU 內(nèi)核。英偉達(dá)表示，其對(duì)于 DeepSeek-R1 的最新進(jìn)展及其應(yīng)用潛力感到非常興奮。