比英偉達(dá)工程師還熟練!DeepSeek R1+測試時Scaling自動優(yōu)化GPU內(nèi)核
本周英偉達(dá)的一篇技術(shù)博客引發(fā)了業(yè)界震動!
英偉達(dá)的團(tuán)隊嘗試?yán)肈eepSeek-R1和推理時擴(kuò)展實現(xiàn)GPU內(nèi)核生成自動化,效果極佳。
隨著AI模型的擴(kuò)展,推理時擴(kuò)展(inference-time scaling),也叫測試時擴(kuò)展(test-time scaling)正閃亮登場。
推理時擴(kuò)展就像是給AI模型配備了一個「智慧錦囊」。當(dāng)模型進(jìn)行推理時,它會額外分配計算資源,讓模型有能力評估多種可能的結(jié)果,然后從中挑選出最優(yōu)解。
通過這種方式,AI可以像人類一樣,有條不紊地剖析復(fù)雜問題,找到最佳解決方案。
為了充分發(fā)揮DeepSeek-R1的優(yōu)勢,同時克服其在生成優(yōu)化GPU內(nèi)核時遇到的困難,英偉達(dá)的工程師們想出了一個創(chuàng)新的方法——將推理時擴(kuò)展技術(shù)與DeepSeek-R1相結(jié)合,構(gòu)建了一種全新的工作流程。
他們使用DeepSeek-R1,在推理過程中借助額外的計算能力來解決一個復(fù)雜問題:旨在自動生成數(shù)值正確且針對不同類型注意力機(jī)制進(jìn)行優(yōu)化的GPU注意力內(nèi)核,整個過程無需任何顯式編程。
在某些情況下,R1生成的內(nèi)核甚至比嫻熟的工程師開發(fā)出來的還要出色!
對此,網(wǎng)友評價道:「英偉達(dá)是在毀掉自己的護(hù)城河嗎?」
優(yōu)化注意力內(nèi)核的挑戰(zhàn)
深入了解推理時擴(kuò)展技術(shù)如何發(fā)揮作用之前,先要認(rèn)識一個LLM中至關(guān)重要的概念——注意力機(jī)制。
注意力機(jī)制就像是「聚光燈」,能讓AI模型在處理任務(wù)時,有選擇地聚焦在輸入信息中最相關(guān)的部分,快速找到關(guān)鍵語句,做出更準(zhǔn)確的預(yù)測,發(fā)現(xiàn)數(shù)據(jù)中的隱藏模式。
但是,注意力操作的計算復(fù)雜度與輸入序列長度的平方成正比。輸入的文本越長,模型處理起來就會越吃力,不僅計算量大幅增加,還可能出現(xiàn)運(yùn)行時錯誤,比如內(nèi)存不足的情況。
為了避免這些問題,提高計算效率,開發(fā)優(yōu)化的底層實現(xiàn),也就是GPU內(nèi)核,十分必要。
另外,注意力機(jī)制有多種不同的變體,像因果注意力、相對位置嵌入、alibi等。工程師們在面對不同的任務(wù)時,往往需要把這些變體組合起來使用。
在多模態(tài)模型,比如視覺Transformer中,需要專門的注意力機(jī)制,像空間鄰域注意力(Spatial Neighborhood Attention),來處理計算機(jī)視覺、視頻生成模型中常見的時空信息。
開發(fā)一個優(yōu)化的GPU注意力內(nèi)核,對經(jīng)驗豐富的軟件工程師來說,也是一項艱巨的任務(wù),費(fèi)時費(fèi)力。
盡管像DeepSeek-R1這樣的模型在代碼生成任務(wù)中展現(xiàn)出了很大的潛力,但它們在第一次嘗試生成優(yōu)化代碼時,還是會遇到不少挑戰(zhàn)。
這使得在推理時必須使用其他策略來生成優(yōu)化代碼。
以下是為相對位置嵌入注意力內(nèi)核輸入的示例用戶提示。
模型有時會產(chǎn)生幻覺,生成一些「不靠譜」的代碼,要么在語法上有錯誤,要么把不同語言或框架的語法混在一起,導(dǎo)致無法運(yùn)行或效率低下。
計算最優(yōu)的GPU線程映射也不是一件容易的事,通常需要反復(fù)調(diào)整優(yōu)化,才能得到一個正確又高效的內(nèi)核。
DeepSeek R1與推理時擴(kuò)展「強(qiáng)強(qiáng)聯(lián)合」
為了攻克優(yōu)化GPU注意力內(nèi)核這個難題,英偉達(dá)的工程師們想出了一個巧妙的辦法,他們把DeepSeek-R1模型和推理時擴(kuò)展技術(shù)結(jié)合起來,創(chuàng)造了一種新的工作流程。
一開始,工程師會手動輸入一個提示。然后,DeepSeek-R1會根據(jù)這個提示,在第一次遍歷中生成GPU代碼,也就是內(nèi)核代碼。
生成的代碼會交給一個特殊的驗證器,這個驗證器運(yùn)行在英偉達(dá)H100 GPU上,仔細(xì)分析生成的內(nèi)核代碼。
如果發(fā)現(xiàn)代碼有不足的地方,驗證器就會生成新的提示,再把這些提示作為輸入,反饋給DeepSeek-R1。模型根據(jù)新的提示,對代碼進(jìn)行改進(jìn),如此循環(huán)往復(fù)。
工程師們發(fā)現(xiàn),這個過程持續(xù)15分鐘,就能得到一個性能更好的注意力內(nèi)核。
根據(jù)斯坦福大學(xué)的KernelBench基準(zhǔn)測試,它生成的內(nèi)核在處理Level-1問題時,在數(shù)值上100%是正確的。在處理Level-2問題時,正確率也能達(dá)到96% 。
KernelBench Level-1問題解決率,是用來評估LLM為特定計算任務(wù)生成高效GPU內(nèi)核能力的數(shù)值正確指標(biāo),它是一系列測試LLM GPU編程能力挑戰(zhàn)的一部分。
推理時間預(yù)算對生成正確內(nèi)核的影響也很明顯。從測試結(jié)果來看,在Level-1問題中,如果每個問題分配的推理時間超過10分鐘,就能為100個問題中的大多數(shù)生成數(shù)值正確的代碼。
這也意味著,給模型足夠的「思考時間」,它真的能給出更好的答案。
利用DeepSeek-R1模型,通過在推理時投入更多計算資源,可以生成比一些熟練工程師開發(fā)的優(yōu)化內(nèi)核還要好的結(jié)果,這為GPU內(nèi)核的自動化生成開辟了一條新的道路。
目前這項技術(shù)還處于早期研究階段,雖然已經(jīng)取得了一些令人振奮的成果,但要想讓它能穩(wěn)定地產(chǎn)生更好的結(jié)果,還有很多工作要做。
研究者對DeepSeek-R1的最新進(jìn)展及潛力感到興奮。