本周英偉達(dá)的一篇技術(shù)博客引發(fā)了業(yè)界震動(dòng)！

英偉達(dá)的團(tuán)隊(duì)嘗試?yán)肈eepSeek-R1和推理時(shí)擴(kuò)展實(shí)現(xiàn)GPU內(nèi)核生成自動(dòng)化，效果極佳。

隨著AI模型的擴(kuò)展，推理時(shí)擴(kuò)展（inference-time scaling），也叫測試時(shí)擴(kuò)展（test-time scaling）正閃亮登場。

推理時(shí)擴(kuò)展就像是給AI模型配備了一個(gè)「智慧錦囊」。當(dāng)模型進(jìn)行推理時(shí)，它會(huì)額外分配計(jì)算資源，讓模型有能力評估多種可能的結(jié)果，然后從中挑選出最優(yōu)解。

通過這種方式，AI可以像人類一樣，有條不紊地剖析復(fù)雜問題，找到最佳解決方案。

為了充分發(fā)揮DeepSeek-R1的優(yōu)勢，同時(shí)克服其在生成優(yōu)化GPU內(nèi)核時(shí)遇到的困難，英偉達(dá)的工程師們想出了一個(gè)創(chuàng)新的方法——將推理時(shí)擴(kuò)展技術(shù)與DeepSeek-R1相結(jié)合，構(gòu)建了一種全新的工作流程。

他們使用DeepSeek-R1，在推理過程中借助額外的計(jì)算能力來解決一個(gè)復(fù)雜問題：旨在自動(dòng)生成數(shù)值正確且針對不同類型注意力機(jī)制進(jìn)行優(yōu)化的GPU注意力內(nèi)核，整個(gè)過程無需任何顯式編程。

在某些情況下，R1生成的內(nèi)核甚至比嫻熟的工程師開發(fā)出來的還要出色！

對此，網(wǎng)友評價(jià)道：「英偉達(dá)是在毀掉自己的護(hù)城河嗎？」

優(yōu)化注意力內(nèi)核的挑戰(zhàn)

深入了解推理時(shí)擴(kuò)展技術(shù)如何發(fā)揮作用之前，先要認(rèn)識(shí)一個(gè)LLM中至關(guān)重要的概念——注意力機(jī)制。

注意力機(jī)制就像是「聚光燈」，能讓AI模型在處理任務(wù)時(shí)，有選擇地聚焦在輸入信息中最相關(guān)的部分，快速找到關(guān)鍵語句，做出更準(zhǔn)確的預(yù)測，發(fā)現(xiàn)數(shù)據(jù)中的隱藏模式。

但是，注意力操作的計(jì)算復(fù)雜度與輸入序列長度的平方成正比。輸入的文本越長，模型處理起來就會(huì)越吃力，不僅計(jì)算量大幅增加，還可能出現(xiàn)運(yùn)行時(shí)錯(cuò)誤，比如內(nèi)存不足的情況。

為了避免這些問題，提高計(jì)算效率，開發(fā)優(yōu)化的底層實(shí)現(xiàn)，也就是GPU內(nèi)核，十分必要。

另外，注意力機(jī)制有多種不同的變體，像因果注意力、相對位置嵌入、alibi等。工程師們在面對不同的任務(wù)時(shí)，往往需要把這些變體組合起來使用。

在多模態(tài)模型，比如視覺Transformer中，需要專門的注意力機(jī)制，像空間鄰域注意力（Spatial Neighborhood Attention），來處理計(jì)算機(jī)視覺、視頻生成模型中常見的時(shí)空信息。

開發(fā)一個(gè)優(yōu)化的GPU注意力內(nèi)核，對經(jīng)驗(yàn)豐富的軟件工程師來說，也是一項(xiàng)艱巨的任務(wù)，費(fèi)時(shí)費(fèi)力。

盡管像DeepSeek-R1這樣的模型在代碼生成任務(wù)中展現(xiàn)出了很大的潛力，但它們在第一次嘗試生成優(yōu)化代碼時(shí)，還是會(huì)遇到不少挑戰(zhàn)。

這使得在推理時(shí)必須使用其他策略來生成優(yōu)化代碼。

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

模型有時(shí)會(huì)產(chǎn)生幻覺，生成一些「不靠譜」的代碼，要么在語法上有錯(cuò)誤，要么把不同語言或框架的語法混在一起，導(dǎo)致無法運(yùn)行或效率低下。

計(jì)算最優(yōu)的GPU線程映射也不是一件容易的事，通常需要反復(fù)調(diào)整優(yōu)化，才能得到一個(gè)正確又高效的內(nèi)核。

DeepSeek R1與推理時(shí)擴(kuò)展「強(qiáng)強(qiáng)聯(lián)合」

為了攻克優(yōu)化GPU注意力內(nèi)核這個(gè)難題，英偉達(dá)的工程師們想出了一個(gè)巧妙的辦法，他們把DeepSeek-R1模型和推理時(shí)擴(kuò)展技術(shù)結(jié)合起來，創(chuàng)造了一種新的工作流程。

一開始，工程師會(huì)手動(dòng)輸入一個(gè)提示。然后，DeepSeek-R1會(huì)根據(jù)這個(gè)提示，在第一次遍歷中生成GPU代碼，也就是內(nèi)核代碼。

生成的代碼會(huì)交給一個(gè)特殊的驗(yàn)證器，這個(gè)驗(yàn)證器運(yùn)行在英偉達(dá)H100 GPU上，仔細(xì)分析生成的內(nèi)核代碼。

如果發(fā)現(xiàn)代碼有不足的地方，驗(yàn)證器就會(huì)生成新的提示，再把這些提示作為輸入，反饋給DeepSeek-R1。模型根據(jù)新的提示，對代碼進(jìn)行改進(jìn)，如此循環(huán)往復(fù)。

工程師們發(fā)現(xiàn)，這個(gè)過程持續(xù)15分鐘，就能得到一個(gè)性能更好的注意力內(nèi)核。

根據(jù)斯坦福大學(xué)的KernelBench基準(zhǔn)測試，它生成的內(nèi)核在處理Level-1問題時(shí)，在數(shù)值上100%是正確的。在處理Level-2問題時(shí)，正確率也能達(dá)到96% 。

KernelBench Level-1問題解決率，是用來評估LLM為特定計(jì)算任務(wù)生成高效GPU內(nèi)核能力的數(shù)值正確指標(biāo)，它是一系列測試LLM GPU編程能力挑戰(zhàn)的一部分。

推理時(shí)間預(yù)算對生成正確內(nèi)核的影響也很明顯。從測試結(jié)果來看，在Level-1問題中，如果每個(gè)問題分配的推理時(shí)間超過10分鐘，就能為100個(gè)問題中的大多數(shù)生成數(shù)值正確的代碼。

這也意味著，給模型足夠的「思考時(shí)間」，它真的能給出更好的答案。

利用DeepSeek-R1模型，通過在推理時(shí)投入更多計(jì)算資源，可以生成比一些熟練工程師開發(fā)的優(yōu)化內(nèi)核還要好的結(jié)果，這為GPU內(nèi)核的自動(dòng)化生成開辟了一條新的道路。

目前這項(xiàng)技術(shù)還處于早期研究階段，雖然已經(jīng)取得了一些令人振奮的成果，但要想讓它能穩(wěn)定地產(chǎn)生更好的結(jié)果，還有很多工作要做。

研究者對DeepSeek-R1的最新進(jìn)展及潛力感到興奮。