存算一體芯片如何支持Transformer等不同模型？

后摩智能致力于打造通用人工智能芯片，自主研發(fā)的存算一體芯片在支持各類模型方面表現(xiàn)突出，包括YOLO系列網(wǎng)絡(luò)、BEV系列網(wǎng)絡(luò)、點(diǎn)云系列網(wǎng)絡(luò)等。這一系列芯片不僅在性能上有著顯著的優(yōu)勢(shì)，而且特別針對(duì)目前自動(dòng)駕駛領(lǐng)域的算法進(jìn)行了專門的優(yōu)化。近期，后摩智能剛完成新一款技術(shù)驗(yàn)證芯片的量產(chǎn)測(cè)試，屬于國(guó)內(nèi)首款基于存算一體架構(gòu)的7nm車規(guī)級(jí)技術(shù)驗(yàn)證芯片，專為Transformer 等車端大模型設(shè)計(jì)。

當(dāng)前，自動(dòng)駕駛等領(lǐng)域中，Transformer模型的應(yīng)用逐漸占據(jù)主導(dǎo)地位。這一架構(gòu)的優(yōu)勢(shì)在于其能夠更好地捕捉長(zhǎng)距離依賴關(guān)系，有助于提高模型對(duì)復(fù)雜場(chǎng)景的理解和處理能力。面對(duì)Transformer在自動(dòng)駕駛中的日益增長(zhǎng)的需求，一個(gè)關(guān)鍵的問(wèn)題浮現(xiàn)出來(lái)：存算一體芯片是否能夠高效部署Transformer模型？

Q1?存算一體芯片能高效部署Transformer嗎？

Transformer架構(gòu)是一種強(qiáng)大的神經(jīng)網(wǎng)絡(luò)架構(gòu)，能夠?qū)崿F(xiàn)高效的序列建模和復(fù)雜的任務(wù)處理，它的核心組成部分包括多層感知機(jī)（MLP）和多頭注意力（MHA）。MLP是一種前饋神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu)，由多個(gè)層次的神經(jīng)元組成，每一層都與上一層的所有神經(jīng)元相連。

這些神經(jīng)元通過(guò)權(quán)重進(jìn)行連接。MLP在Transformer中負(fù)責(zé)對(duì)輸入特征進(jìn)行變換和映射，幫助網(wǎng)絡(luò)捕捉不同層次的抽象特征；MHA允許網(wǎng)絡(luò)在不同位置對(duì)輸入序列的不同部分進(jìn)行關(guān)注，從而提高模型的并行性和全局信息的捕捉能力。MHA的基本思想是通過(guò)多個(gè)注意力頭（Attention Head）并行處理輸入序列，每個(gè)頭都學(xué)習(xí)關(guān)注輸入的不同方面。

這種并行性有助于有效地處理長(zhǎng)序列，并使網(wǎng)絡(luò)更具擴(kuò)展性和泛化能力。在部署Transformer時(shí)，AI芯片的任務(wù)就是高效執(zhí)行網(wǎng)絡(luò)中的MLP和MHA結(jié)構(gòu)。這需要AI芯片能對(duì)其中所有算子都有很高的并行執(zhí)行能力。

Q2?多層感知機(jī)如何高效部署在存算一體芯片上？

全連接層的本質(zhì)是執(zhí)行兩個(gè)矩陣（輸入特征矩陣和權(quán)重矩陣）的矩陣乘法。我們將其中權(quán)重存放在存算單元上。將輸入的特征送入存算單元，由存算單元中的計(jì)算單元直接完成乘累加操作，輸出乘累加結(jié)果。由于存算單元的計(jì)算密度很高，執(zhí)行這種全連接層的并行度很高，因此效率很高。另一方面，權(quán)重被保持在存算單元上，不發(fā)生移動(dòng)，從而大幅降低了搬移權(quán)重所帶來(lái)的能耗開(kāi)銷。

多頭注意力是Transformer模型的一個(gè)關(guān)鍵組成部分，它有點(diǎn)像大腦的多個(gè)小模塊，每個(gè)模塊都負(fù)責(zé)關(guān)注輸入數(shù)據(jù)的不同方面。這個(gè)結(jié)構(gòu)之所以特別，是因?yàn)樗硕鄠€(gè)小“頭”，每個(gè)“頭”都在關(guān)注輸入數(shù)據(jù)的不同部分。

在處理輸入數(shù)據(jù)時(shí)，每個(gè)“頭”都有三個(gè)關(guān)鍵矩陣，分別是查詢（Q）、鍵（K）、和值（V）。通過(guò)一系列數(shù)學(xué)運(yùn)算，多頭注意力可以捕捉到輸入序列中不同位置之間的關(guān)系。

關(guān)系計(jì)算：首先，我們讓查詢（Q）和鍵（K）進(jìn)行一種特殊的數(shù)學(xué)操作，就像在查找輸入數(shù)據(jù)中不同部分之間的聯(lián)系。這為模型提供了對(duì)輸入序列中不同位置的關(guān)注程度。

重要性映射：接著，我們通過(guò)一個(gè)函數(shù)（softmax）把剛才計(jì)算的結(jié)果映射到0到1之間，就好像在給不同位置分配注意力的“權(quán)重”，表示它們的相對(duì)重要性。

信息整合：最后，我們把剛才得到的歸一化的結(jié)果與值（V）進(jìn)行另一次數(shù)學(xué)操作，這樣就得到了最終輸出。這一步把被注意到的值通過(guò)權(quán)重相加，得到多頭注意力的最終輸出。

雖然這里提到的數(shù)學(xué)操作和全連接層有點(diǎn)相似，但在多頭注意力中，查詢、鍵、和值這三個(gè)矩陣是動(dòng)態(tài)生成的。這就意味著在執(zhí)行數(shù)學(xué)操作時(shí)，需要靈活的加載數(shù)據(jù)到存算單元上，這一過(guò)程的效率對(duì)于處理器性能非常關(guān)鍵。

為了解決這個(gè)問(wèn)題，后摩智能設(shè)計(jì)了高效的存算單元數(shù)據(jù)加載硬件。這可以極大地提高存算單元中數(shù)據(jù)的替換效率，確保在多頭注意力的計(jì)算中，動(dòng)態(tài)產(chǎn)生的矩陣K、V能夠快速而高效地加載到存算單元中。這種巧妙的設(shè)計(jì)使得存算一體芯片能夠在執(zhí)行多頭注意力結(jié)構(gòu)時(shí)取得最佳性能，為Transformer模型的高效運(yùn)行提供了強(qiáng)有力的支持。

除了計(jì)算密集型的全連接層和矩陣乘法之外，后摩智能的芯片還擁有大量的向量算力和標(biāo)量算力來(lái)處理其它算子，例如softmax和layernorm算子。為了充分利用這些算力資源，后摩智能采用了先進(jìn)的編譯優(yōu)化算法。這一算法的設(shè)計(jì)使得存算單元、向量單元和標(biāo)量單元能夠被同時(shí)調(diào)度，實(shí)現(xiàn)并行執(zhí)行不同的計(jì)算任務(wù)。這種并行計(jì)算的優(yōu)勢(shì)不僅僅體現(xiàn)在同一算子的多個(gè)實(shí)例之間，更在于不同算子之間的并發(fā)執(zhí)行。通過(guò)同時(shí)處理各個(gè)算子，后摩智能的芯片在運(yùn)行Transformer時(shí)能夠達(dá)到最高的效率，極大地提升了整個(gè)計(jì)算過(guò)程的速度和效能。

綜合而言，后摩智能芯片以其先進(jìn)的設(shè)計(jì)理念和高效的存算一體架構(gòu)，成功解決了對(duì)Transformer算法的高效支持問(wèn)題。通過(guò)優(yōu)化全連接層、矩陣乘法和動(dòng)態(tài)加載、向量和標(biāo)量運(yùn)算等關(guān)鍵環(huán)節(jié)，后摩智能為Transformer運(yùn)行的高效性和性能提供了強(qiáng)大的支持，為自動(dòng)駕駛等領(lǐng)域的發(fā)展注入了新的動(dòng)力。未來(lái)，后摩智能的存算一體架構(gòu)芯片也將隨著算法與硬件的演進(jìn)，不斷迭代，滿足萬(wàn)物智能時(shí)代對(duì)算力的澎湃需求。

審核編輯：劉清