1. 從算力饑渴到專(zhuān)用計(jì)算AI芯片這十年的三條岔路過(guò)去十年AI圈子里有個(gè)很有趣的“剪刀差”做算法的同學(xué)每年都在抱怨模型越來(lái)越大、訓(xùn)練越來(lái)越慢而做硬件的同學(xué)則在拼命堆晶體管、提頻率。兩邊都很努力但總感覺(jué)隔著一條河。直到我真正跳進(jìn)AI芯片這個(gè)坑才明白問(wèn)題不在“算力不夠”而在“算力用不對(duì)地方”。2012年AlexNet在GPU上跑出驚人效果之后行業(yè)其實(shí)面臨一個(gè)岔路口。第一條路是把GPU這條通用并行計(jì)算路線(xiàn)做到極致靠CUDA生態(tài)和Tensor Core持續(xù)迭代第二條路是走FPGA用可重構(gòu)邏輯去適配快速變化的算法第三條路則是押注ASIC直接為神經(jīng)網(wǎng)絡(luò)“刻”一塊專(zhuān)用硅片。當(dāng)時(shí)誰(shuí)也不敢說(shuō)哪條路一定贏但十年后回頭看三條路都活了下來(lái)只是各自找到了完全不同的生存空間。這背后的邏輯并不復(fù)雜。AI計(jì)算和傳統(tǒng)計(jì)算最大的區(qū)別在于傳統(tǒng)計(jì)算是“控制密集”比如操作系統(tǒng)、數(shù)據(jù)庫(kù)這類(lèi)負(fù)載大量的時(shí)間花在分支判斷、地址跳轉(zhuǎn)、中斷處理上需要的是強(qiáng)大的控制邏輯和通用性而AI計(jì)算是“數(shù)據(jù)密集”卷積、矩陣乘、激活函數(shù)、歸一化這些操作結(jié)構(gòu)高度規(guī)整數(shù)據(jù)可以流水線(xiàn)式地灌進(jìn)去幾乎不需要復(fù)雜的控制流。一旦負(fù)載形態(tài)變成了這樣通用CPU那種“每一條指令都要先取指、解碼、再執(zhí)行”的流水線(xiàn)方式就顯得笨重了——大量晶體管和功耗都消耗在“準(zhǔn)備執(zhí)行”這個(gè)過(guò)程里而不是真正“執(zhí)行”上。所以AI芯片架構(gòu)設(shè)計(jì)的核心命題從第一天起就不是“怎么把芯片做快”而是“怎么把晶體管花在刀刃上”。這個(gè)“刀刃”就是計(jì)算單元本身。圍繞這個(gè)命題不同團(tuán)隊(duì)給出了截然不同的答案。本文不打算寫(xiě)成教科書(shū)式的綜述只是想把我在實(shí)際接觸和評(píng)估各種AI芯片架構(gòu)時(shí)看到的門(mén)道、走過(guò)的彎路、以及那些在PPT上看不到的真實(shí)權(quán)衡系統(tǒng)地?cái)傞_(kāi)來(lái)聊一聊。適合正在做AI基礎(chǔ)設(shè)施選型、做算法工程化落地、或者想從軟件側(cè)理解硬件邏輯的讀者。2. 計(jì)算陣列、存儲(chǔ)墻與片上互聯(lián)拆開(kāi)一塊AI芯片的三個(gè)橫截面拿到任何一款A(yù)I芯片的架構(gòu)圖無(wú)論是NVIDIA的Hopper、Google的TPU還是國(guó)內(nèi)廠商的各類(lèi)NPU第一眼看上去都會(huì)覺(jué)得眼花繚亂。但如果拆開(kāi)看其實(shí)所有AI芯片都在回答同一個(gè)問(wèn)題的三個(gè)子問(wèn)題計(jì)算怎么做、數(shù)據(jù)怎么喂、數(shù)據(jù)怎么搬。2.1 計(jì)算陣列從二維到三維的MAC組織方式計(jì)算單元是AI芯片的心臟而心臟的基本跳動(dòng)單元叫MAC也就是乘累加操作。一次MAC做的事情很簡(jiǎn)單兩個(gè)數(shù)相乘再加到累加器上。但神經(jīng)網(wǎng)絡(luò)里的矩陣乘法動(dòng)輒上百億次MAC怎么把這些MAC組織起來(lái)直接決定了芯片的面積效率、功耗效率和時(shí)鐘頻率。主流的組織方式有三種。第一種是二維脈動(dòng)陣列TPU是典型代表。數(shù)據(jù)像波浪一樣在二維陣列里“流動(dòng)”每個(gè)計(jì)算單元只和相鄰單元通信不需要全局廣播。這種做法的好處是數(shù)據(jù)復(fù)用率極高、布線(xiàn)短、功耗低壞處是陣列利用率受矩陣形狀影響很大——如果矩陣維度和陣列尺寸不匹配會(huì)有大量計(jì)算單元空轉(zhuǎn)。第二種是SIMD風(fēng)格的寬向量單元GPU的Tensor Core和我見(jiàn)過(guò)的一些NPU都走這條路。數(shù)據(jù)被切成固定寬度的向量比如16x16、32x32的tile一條指令驅(qū)動(dòng)整個(gè)陣列同步執(zhí)行。這種設(shè)計(jì)對(duì)編譯器非常友好利用率容易做高但對(duì)數(shù)據(jù)在片上存儲(chǔ)器的布局要求苛刻排不好就會(huì)出現(xiàn)大量bank沖突。第三種是近兩年開(kāi)始流行的一些分布式/可重構(gòu)陣列計(jì)算單元之間通過(guò)片上網(wǎng)絡(luò)動(dòng)態(tài)連接理論上可以靈活映射任意結(jié)構(gòu)的計(jì)算圖實(shí)際落地難度確實(shí)也比較大。我在實(shí)際評(píng)估芯片時(shí)很少只看峰值算力更關(guān)心“有效算力密度”——也就是在給定芯片面積和功耗下持續(xù)跑到70%以上利用率的算力是多少。這個(gè)指標(biāo)才真實(shí)反映MAC組織方式的優(yōu)劣。2.2 存儲(chǔ)層級(jí)為什么說(shuō)“存儲(chǔ)墻”是AI芯片的第一瓶頸很多芯片的Datasheet上寫(xiě)著幾十甚至上百TOPS的算力跑個(gè)小模型也確實(shí)能跑到接近理論峰值。但一上真實(shí)模型性能直接腰斬甚至膝斬。絕大多數(shù)情況下問(wèn)題不是出在計(jì)算單元上而是出在存儲(chǔ)上。AI芯片的存儲(chǔ)層級(jí)大致分四層HBM或者LPDDR這樣的片外DRAM、片上SRAM有些廠商叫Shared Memory、有些叫Buffer、寄存器堆、以及計(jì)算單元內(nèi)部的累加器/暫存器。帶寬從上到下是數(shù)量級(jí)的提升容量則是數(shù)量級(jí)的下降。真實(shí)計(jì)算中要跑一個(gè)矩陣乘你得先把A矩陣的塊和B矩陣的塊從DRAM搬到SRAM算完之后把結(jié)果C搬回去然后再搬下一塊。如果搬數(shù)據(jù)的時(shí)間比計(jì)算時(shí)間還長(zhǎng)計(jì)算單元就只能閑著等數(shù)據(jù)。這就是“存儲(chǔ)墻”的本質(zhì)算力漲得比存儲(chǔ)帶寬快得多導(dǎo)致喂數(shù)據(jù)的速度跟不上消耗數(shù)據(jù)的速度。解決思路無(wú)外乎三條路一是把數(shù)據(jù)盡量留在片上通過(guò)tiling切分讓數(shù)據(jù)反復(fù)復(fù)用典型如CUTLASS、triton這類(lèi)庫(kù)在做的事二是設(shè)計(jì)更大的片上SRAM比如有些芯片把SRAM做到幾十兆甚至上百兆目的就是能放下更大的tile三是做近存計(jì)算或存內(nèi)計(jì)算把計(jì)算邏輯直接懟到存儲(chǔ)單元旁邊甚至里面這是架構(gòu)層面的終極解法后面會(huì)專(zhuān)門(mén)講。從軟件側(cè)來(lái)看存儲(chǔ)層級(jí)設(shè)計(jì)直接決定了算子的實(shí)現(xiàn)策略。同樣一個(gè)卷積在A芯片上要分成四段來(lái)搬數(shù)據(jù)在B芯片上也許一次就能裝下整張feature map。編程模型再抽象最終都是要映射到物理的存儲(chǔ)層級(jí)上的。所以評(píng)估一款A(yù)I芯片我一定先要拿到它片上SRAM的容量和帶寬數(shù)據(jù)再去看峰值算力——順序不能反。2.3 片上互聯(lián)多核架構(gòu)最容易忽略的暗坑現(xiàn)在的AI芯片沒(méi)有單核打天下的。即便單die內(nèi)部也會(huì)拆分成多個(gè)計(jì)算核比如NVIDIA的SM、Habana的TPC、各家NPU的NPU Core多核之間怎么通信是芯片能不能“人多吃上飯”的分水嶺。ARM架構(gòu)的SoC一般用AMBA總線(xiàn)簡(jiǎn)單直接適合小規(guī)模互聯(lián)。AI芯片上常見(jiàn)的互聯(lián)方式有三種網(wǎng)格片上網(wǎng)絡(luò)、環(huán)形總線(xiàn)、以及高帶寬的crossbar。網(wǎng)格片上網(wǎng)絡(luò)代表是某些云端AI芯片的優(yōu)點(diǎn)是可擴(kuò)展性好核多了也能掛得住代價(jià)是通信延遲不可控?cái)?shù)據(jù)從一個(gè)核到另一個(gè)核要跳很多跳編譯器必須做精確的通信編排否則性能穩(wěn)定性很難保證。環(huán)形總線(xiàn)NVIDIA GPU的NVLink片上版本可以近似這么理解還有不少ASIC沿用實(shí)現(xiàn)簡(jiǎn)單廣播效率高但環(huán)一旦做大了單跳帶寬和整體延遲會(huì)惡化。Crossbar是最奢侈的方案任意核到任意核都是直達(dá)延遲低、編程簡(jiǎn)單但面積和布線(xiàn)開(kāi)銷(xiāo)隨著核數(shù)平方級(jí)上漲超過(guò)一定規(guī)模就完全不經(jīng)濟(jì)了。我見(jiàn)過(guò)不少項(xiàng)目在選型時(shí)只看每個(gè)核的算力忽略了核間通信能力。結(jié)果就是單核跑一個(gè)算子很快但一旦模型是并行度要求高的結(jié)構(gòu)通信帶寬就成了瓶頸。這個(gè)坑在2D卷積這類(lèi)本來(lái)就有空間局部性的算子里不太明顯但在Transformer的注意力計(jì)算里會(huì)非常顯著——因?yàn)镼、K、V矩陣要頻繁做跨頭的拼接和切分片上數(shù)據(jù)搬運(yùn)量極大。3. 指令集架構(gòu)上行那些繞不開(kāi)的SIMD、VLIW與可重構(gòu)軟件開(kāi)發(fā)者說(shuō)“架構(gòu)”腦子里浮現(xiàn)的往往是操作系統(tǒng)的分層、微服務(wù)怎么拆、前后端怎么通信。但芯片領(lǐng)域的“架構(gòu)”最初的落腳點(diǎn)是指令集架構(gòu)——也就是軟件和硬件之間的那條合同。AI芯片的指令集設(shè)計(jì)其實(shí)暴露了芯片公司對(duì)“軟件棧應(yīng)該長(zhǎng)什么樣”的基本立場(chǎng)。3.1 SIMD與向量化GPU走通的那條康莊大道SIMD不是新概念上世紀(jì)六十年代就有向量機(jī)的雛形但把它在AI時(shí)代發(fā)揚(yáng)光大的是GPU。GPU的哲學(xué)是“大量簡(jiǎn)單計(jì)算單元聽(tīng)同一個(gè)指揮做同樣的事”。每個(gè)線(xiàn)程執(zhí)行同一條指令但操作不同的數(shù)據(jù)。這個(gè)模式下編譯器或者程序員要操心的問(wèn)題只有一個(gè)把數(shù)據(jù)打包成合適的向量寬度確保所有線(xiàn)程都在工作。放到AI芯片的語(yǔ)境下SIMD風(fēng)格最典型的代表就是Tensor Core和各類(lèi)NPU的MAC陣列指令。比如英偉達(dá)的wmma或者mma指令一次可以算一個(gè)16x16x16的矩陣乘。如果你寫(xiě)CUDA寫(xiě)過(guò)這種指令你一定會(huì)感覺(jué)到它本質(zhì)上是在“教”硬件怎么切分矩陣、怎么排列數(shù)據(jù)。這種模式的優(yōu)點(diǎn)是想清楚之后效率可以做到極高缺點(diǎn)則是軟件棧的復(fù)雜度被轉(zhuǎn)移給了開(kāi)發(fā)者——如何把一個(gè)大矩陣乘拆成無(wú)數(shù)個(gè)tile并且保證每個(gè)tile的加載和計(jì)算重疊得嚴(yán)絲合縫這不是一個(gè)輕松的活。3.2 VLIW的誘惑與陷阱編譯器把命脈握在自己手里的野心VLIW和SIMD容易混淆但思路完全不同。SIMD是“一條指令操作很多數(shù)據(jù)”VLIW是“一條指令里有好幾個(gè)操作并行發(fā)射到不同功能單元”。VLIW把亂序執(zhí)行、動(dòng)態(tài)調(diào)度這些讓CPU頭疼的硬件邏輯全部搬到了編譯期。編譯器靜態(tài)分析代碼把能并行的操作打包成一條很長(zhǎng)的指令硬件只需要無(wú)腦執(zhí)行。這套哲學(xué)在DSP領(lǐng)域統(tǒng)治了很久在AI芯片領(lǐng)域也曾被寄予厚望。誘惑在于它省掉了大量亂序執(zhí)行的硬件面積和功耗都能降下來(lái)對(duì)追求極致能效比的推理芯片非常有吸引力。陷阱則在于編譯器需要完美預(yù)測(cè)運(yùn)行時(shí)的情況一旦遇到緩存未命中、分支跳轉(zhuǎn)這類(lèi)不確定事件指令流水線(xiàn)就會(huì)氣泡叢生性能大幅回退。AI計(jì)算里最典型的就是動(dòng)態(tài)shape——序列長(zhǎng)度不是固定的導(dǎo)致循環(huán)次數(shù)、內(nèi)存訪問(wèn)模式全部動(dòng)態(tài)化。VLIW編譯器的靜態(tài)分析在這類(lèi)場(chǎng)景下非常痛苦。我實(shí)際接觸過(guò)的態(tài)度是純VLIW適合算法高度固定的專(zhuān)用芯片比如某個(gè)特定檢測(cè)模型、某個(gè)固定尺寸的推薦模型。一旦要做多模型、動(dòng)態(tài)shape、復(fù)雜控制流的通用AI推理VLIW的優(yōu)勢(shì)會(huì)迅速被編譯器的力不從心吞噬。3.3 RISC-V對(duì)AI芯片的影響不只是“開(kāi)源”兩個(gè)字可以概括這幾年幾乎每個(gè)做AI芯片的創(chuàng)業(yè)公司都會(huì)在PPT里提一句“基于RISC-V”。但RISC-V對(duì)AI芯片的意義并不是“免費(fèi)用一個(gè)指令集”而是“可以干凈地?cái)U(kuò)展自定義指令”。傳統(tǒng)CPU有x86和ARM兩座大山指令集僵化想加一條AI加速指令需要走復(fù)雜的生態(tài)流程。RISC-V則劃出了一大片自定義指令空間芯片公司可以把自己的MAC操作、激活函數(shù)、甚至專(zhuān)用的數(shù)據(jù)搬運(yùn)指令直接做成ISA的一部分。這帶來(lái)的直接收益是編譯器、匯編器和調(diào)試工具可以無(wú)縫支持這些新增指令而不是靠“inline asm”或者“intrinsic函數(shù)”在標(biāo)準(zhǔn)ISA里打補(bǔ)丁。另外一個(gè)容易被忽視的點(diǎn)是RISC-V的向量擴(kuò)展讓AI芯片可以更優(yōu)雅地處理“非矩陣”計(jì)算。神經(jīng)網(wǎng)絡(luò)里除了矩陣乘還有大量element-wise操作比如LayerNorm里的均值方差、Softmax里的指數(shù)歸一化、注意力里的mask操作。這些操作用MAC陣列跑是殺雞用牛刀靠通用CPU核跑又太慢。RISC-V向量單元正好補(bǔ)上這個(gè)空白。所以你會(huì)看到很多AI芯片采用“通用RISC-V核專(zhuān)用NPU陣列”的異構(gòu)設(shè)計(jì)前者管控制流和element-wise后者管重計(jì)算各司其職。4. 從CNN到Transformer神經(jīng)網(wǎng)絡(luò)的演進(jìn)如何反向重塑芯片架構(gòu)AI芯片領(lǐng)域有個(gè)很有意思的現(xiàn)象大多數(shù)芯片架構(gòu)的生命周期比它要跑的模型的生命周期長(zhǎng)很多。一款芯片從設(shè)計(jì)到流片要兩三年而算法社區(qū)半年就能把主流模型換一遍。這意味著芯片架構(gòu)在定義的時(shí)候必須對(duì)“未來(lái)兩三年跑的模型”有足夠強(qiáng)的預(yù)判。4.1 傳統(tǒng)CNN邏輯架構(gòu)在芯片上的映射與優(yōu)化CNN時(shí)代的AI芯片設(shè)計(jì)是相對(duì)“安逸”的因?yàn)榫矸e的結(jié)構(gòu)非常規(guī)整一個(gè)卷積核在輸入特征圖上滑動(dòng)數(shù)據(jù)的空間局部性極好。芯片設(shè)計(jì)者不需要太多創(chuàng)新就能通過(guò)脈動(dòng)陣列、數(shù)據(jù)復(fù)用、行緩沖器這些經(jīng)典手段把卷積跑得很高效。但CNN內(nèi)部其實(shí)也已經(jīng)分化出不同的架構(gòu)流派。以YOLOv8為代表的目標(biāo)檢測(cè)網(wǎng)絡(luò)主干網(wǎng)絡(luò)是卷積、頸部是FPN/PANet做多尺度特征融合、頭部是解耦檢測(cè)頭。這類(lèi)結(jié)構(gòu)在芯片上的映射挑戰(zhàn)在于多尺度特征圖帶來(lái)的內(nèi)存訪問(wèn)不規(guī)律——不同層的特征圖尺寸不同數(shù)據(jù)布局差異大對(duì)存儲(chǔ)帶寬的消耗非?！八槠?。另一類(lèi)是以MobileNet為代表的深度可分離卷積計(jì)算密度低內(nèi)存訪問(wèn)占比高對(duì)MAC陣列的壓力小但對(duì)片上存儲(chǔ)帶寬的壓力反而更大。所以在我做芯片評(píng)估的時(shí)候CNN時(shí)代有個(gè)經(jīng)驗(yàn)法則不要只看理論TOPS要看“有效內(nèi)存帶寬”。一個(gè)TOPS很高但內(nèi)存帶寬捉襟見(jiàn)肘的芯片跑MobileNet這類(lèi)輕量模型可能還不如一個(gè)TOPS稍低但帶寬充裕的芯片。這個(gè)經(jīng)驗(yàn)在Transformer時(shí)代更是被放大到極致。4.2 Transformer架構(gòu)及其工作原理帶來(lái)的三重計(jì)算特征Transformer架構(gòu)及其工作原理對(duì)芯片的沖擊是架構(gòu)層面的。第一重沖擊是注意力機(jī)制的內(nèi)存復(fù)雜度。Q和K矩陣做點(diǎn)積得到注意力分?jǐn)?shù)時(shí)需要產(chǎn)生一個(gè)序列長(zhǎng)度平方級(jí)的中間矩陣。序列長(zhǎng)度從CNN時(shí)代的幾百暴漲到了Transformer時(shí)代的幾千甚至幾萬(wàn)這個(gè)中間矩陣的尺寸直接爆炸——它在片上SRAM里根本放不下必須被頻繁寫(xiě)回DRAM。計(jì)算本身的FLOPs反而還好但數(shù)據(jù)搬運(yùn)的量級(jí)上升了一個(gè)維度。第二重沖擊是動(dòng)態(tài)shape。語(yǔ)言模型的序列長(zhǎng)度是不固定的推理的時(shí)候用戶(hù)輸入的token數(shù)量每次都不同。這讓很多針對(duì)固定shape做過(guò)極致優(yōu)化的靜態(tài)編譯方案失效。我測(cè)試過(guò)一些機(jī)器學(xué)習(xí)編譯器跑固定尺寸的CNN模型可以做到非常接近理論峰值但一旦切換成動(dòng)態(tài)sequence length的GPT類(lèi)模型性能掉一半是常事。原因就在于編譯期無(wú)法預(yù)知循環(huán)邊界指令流水線(xiàn)和數(shù)據(jù)預(yù)取策略全部被打亂。第三重沖擊是計(jì)算和訪存的分離傾向。Transformer里真正“算得多”的是線(xiàn)性層也就是矩陣乘但真正“影響性能”的往往是那些計(jì)算量不大、卻要反復(fù)來(lái)訪存的算子比如Softmax、LayerNorm、Residual Add、以及各種permute和reshape操作。如果芯片的架構(gòu)把MAC陣列做得極其強(qiáng)大卻在element-wise算子和數(shù)據(jù)搬運(yùn)動(dòng)作上偷工減料跑Transformer的時(shí)候整個(gè)流水線(xiàn)會(huì)嚴(yán)重失衡。4.3 為T(mén)ransformer而生的新架構(gòu)方向FlashAttention思路的硬件化軟件層面有一個(gè)很好的解決思路正在快速硬件化那就是FlashAttention的思路。它的核心洞察是不要讓注意力分?jǐn)?shù)矩陣完整地落到DRAM里而是通過(guò)分塊計(jì)算、在線(xiàn)softmax、增量更新讓盡可能多的計(jì)算在SRAM內(nèi)完成。這個(gè)概念看似是一個(gè)算法層面的優(yōu)化實(shí)際上它對(duì)硬件架構(gòu)提出了相當(dāng)具體的訴求片上SRAM要大、計(jì)算單元和SRAM之間的帶寬要足、并且計(jì)算陣列要能高效地流水線(xiàn)化執(zhí)行element-wise操作。我做架構(gòu)評(píng)估時(shí)特別關(guān)注一款芯片對(duì)FlashAttention風(fēng)格算子的支持程度。有些芯片的軟件棧直接提供了融合的Attention算子調(diào)起來(lái)很方便性能也穩(wěn)定有些芯片則要求用戶(hù)自己去手工tile和編排雖然有靈活性但門(mén)檻確實(shí)偏高。這個(gè)差別不僅僅是工程實(shí)現(xiàn)的問(wèn)題本質(zhì)上反映的是芯片架構(gòu)對(duì)數(shù)據(jù)流形態(tài)的適應(yīng)能力——是那種“為矩陣乘而生、其他一概不管”的架構(gòu)還是“理解模型計(jì)算模式、主動(dòng)做算子融合”的架構(gòu)。5. 分布式訓(xùn)練與推理當(dāng)算力從一顆芯片走向一片集群芯片單點(diǎn)的性能再?gòu)?qiáng)在千億參數(shù)模型面前都是杯水車(chē)薪。真正的系統(tǒng)級(jí)問(wèn)題是讓幾千顆芯片協(xié)同工作像一個(gè)整體一樣完成訓(xùn)練或者推理任務(wù)。這個(gè)層面的架構(gòu)設(shè)計(jì)考慮的問(wèn)題從晶體管轉(zhuǎn)移到了機(jī)架、網(wǎng)絡(luò)交換機(jī)和數(shù)據(jù)中心。5.1 三種并行策略背后的通信拓?fù)湫枨蟠竽Ｐ陀?xùn)練的基本功是三種并行數(shù)據(jù)并行、張量并行、流水線(xiàn)并行。三種并行對(duì)芯片間通信的需求截然不同。數(shù)據(jù)并行最簡(jiǎn)單每個(gè)芯片持有完整模型的副本喂不同的數(shù)據(jù)定期同步梯度。同步梯度的通信量取決于模型大小但通信是周期性的、突發(fā)性的可以靠帶寬換延遲。張量并行則是把一層算子的計(jì)算切分成多份讓多顆芯片各算一部分然后通過(guò)AllReduce合并部分結(jié)果。這種并行的通信頻率極高幾乎每做一次矩陣乘就要同步一次對(duì)芯片間互連的延遲和帶寬都是極限壓力。流水線(xiàn)并行稍微溫和一些是把不同層放在不同芯片上數(shù)據(jù)像流水線(xiàn)一樣一層一層往后流通信頻率低但需要足夠大的互連緩沖來(lái)吸收不同層執(zhí)行速度的差異。這三者并不是互斥關(guān)系真實(shí)訓(xùn)練里往往是混合使用的。而混合并行策略對(duì)芯片集群的互連拓?fù)涮岢隽艘粋€(gè)硬性要求不能假設(shè)通信只會(huì)發(fā)生在相鄰節(jié)點(diǎn)之間。所以現(xiàn)代AI集群的互連架構(gòu)基本都采用胖樹(shù)或全連接變體配合RDMA網(wǎng)卡和RoCEv2或者InfiniBand網(wǎng)絡(luò)。芯片本身的片上互聯(lián)能力反而退居其次跨芯片的通信質(zhì)量被網(wǎng)絡(luò)架構(gòu)主導(dǎo)了。5.2 緊耦合與松耦合AI芯片互連架構(gòu)的兩條路線(xiàn)在芯片層面多芯片互連也有兩種路線(xiàn)之爭(zhēng)。一種是緊耦合的典型如NVIDIA的NVLink/NVSwitch方案、NVIDIA的GB200 NVL72這種整機(jī)柜級(jí)互連方案。做到極致的情況下多顆芯片在軟件看來(lái)就像一顆巨大的芯片共享內(nèi)存、統(tǒng)一尋址通信對(duì)上層基本透明。這種方案開(kāi)發(fā)效率極高也是目前大模型訓(xùn)練最順暢的路徑。另一種是松耦合的芯片之間走標(biāo)準(zhǔn)以太網(wǎng)或者定制的高速通信接口每顆芯片獨(dú)立運(yùn)行通過(guò)軟件層的通信庫(kù)類(lèi)似NCCL的機(jī)制來(lái)同步數(shù)據(jù)。這種方案的好處是成本低、靈活性強(qiáng)壞處是通信延遲高、編程復(fù)雜。近年業(yè)界有一些專(zhuān)門(mén)為大模型設(shè)計(jì)的定制化以太網(wǎng)方案例如RoCEv2的擁塞控制算法針對(duì)多對(duì)一通信也就是AllReduce中最容易出現(xiàn)的流量模式做了很多優(yōu)化實(shí)際效果已經(jīng)接近專(zhuān)有網(wǎng)絡(luò)。我在評(píng)估AI芯片集群方案時(shí)會(huì)特別關(guān)注芯片是不是原生支持PCIe點(diǎn)到點(diǎn)通信和NVMe over Fabric這類(lèi)高速互訪協(xié)議。很多芯片單卡跑分很漂亮但互聯(lián)能力孱弱做數(shù)據(jù)并行訓(xùn)練時(shí)效率下降得厲害。芯片之間的通信帶寬、延遲、擁塞控制能力是決定訓(xùn)練集群擴(kuò)展性的隱形天花板。5.3 分布式框架的架構(gòu)層級(jí)從PyTorch DDP到Megatron-DeepSpeed的演進(jìn)分布式訓(xùn)練框架的架構(gòu)層級(jí)演變也值得聊一聊。PyTorch的DDP模式只是實(shí)現(xiàn)了數(shù)據(jù)并行通信邏輯簡(jiǎn)單粗暴每個(gè)step結(jié)束AllReduce同步一次梯度。這個(gè)方案在模型不大、顯卡數(shù)量不夸張的情況下夠用但一旦模型大到單卡裝不下就需要引入張量并行和流水線(xiàn)并行。Megatron-LM提供了完整的張量并行實(shí)現(xiàn)DeepSpeed加入ZeRO零冗余優(yōu)化器來(lái)切分優(yōu)化器狀態(tài)和梯度兩者結(jié)合一般叫Megatron-DeepSpeed就成了大模型訓(xùn)練的標(biāo)配。但對(duì)于AI芯片架構(gòu)師來(lái)說(shuō)這套軟件棧其實(shí)隱含了一層硬件假設(shè)它假設(shè)芯片之間的通信方式是以NCCL為底座的、類(lèi)似GPU的通信語(yǔ)義。一旦要適配新的AI芯片就必須實(shí)現(xiàn)一套類(lèi)似NCCL的通信庫(kù)并且針對(duì)新的互連拓?fù)渥鲂阅苷{(diào)優(yōu)。很多做芯片的公司低估了這部分工作量芯片做出來(lái)之后通信庫(kù)接不上分布式訓(xùn)練根本跑不動(dòng)。所以AI芯片選型時(shí)一定要看它的分布式通信棧成熟度而不光是通信硬件的理論指標(biāo)。6. 架構(gòu)的未來(lái)存算一體、chiplet與軟件定義硬件聊完已經(jīng)落地的方案再往前看一步。AI芯片架構(gòu)的演進(jìn)遠(yuǎn)沒(méi)有到收斂的階段反而因?yàn)槟Ｐ鸵?guī)模的膨脹而進(jìn)入了新的爆發(fā)期。三個(gè)方向在我看來(lái)值得特別關(guān)注。6.1 存算一體繞過(guò)存儲(chǔ)墻的終極手段還是工程烏托邦存算一體Computing-in-Memory的思路非常誘人既然數(shù)據(jù)搬運(yùn)是瓶頸那就不搬了。把計(jì)算單元做成和存儲(chǔ)單元貼合的結(jié)構(gòu)讓數(shù)據(jù)在存儲(chǔ)里就被直接算掉。阻變存儲(chǔ)器RRAM、相變存儲(chǔ)器PCM、甚至靜態(tài)隨機(jī)存儲(chǔ)器SRAM都能做存內(nèi)計(jì)算乘累加操作可以在存儲(chǔ)陣列的字線(xiàn)Word Line和位線(xiàn)Bit Line上直接完成模擬域或數(shù)字域都有實(shí)現(xiàn)案例。但這個(gè)方向的工程挑戰(zhàn)極大。模擬域的存內(nèi)計(jì)算精度受限8bit以上的定點(diǎn)計(jì)算信噪比就很難保證對(duì)主流AI模型尤其是訓(xùn)練場(chǎng)景對(duì)精度要求是FP16起步根本不適用。數(shù)字域的存內(nèi)計(jì)算面積開(kāi)銷(xiāo)又上去了SRAM存內(nèi)計(jì)算的密度相比專(zhuān)用計(jì)算陣列并無(wú)數(shù)量級(jí)優(yōu)勢(shì)。我個(gè)人的判斷是存算一體在未來(lái)幾年很難成為云端訓(xùn)練芯片的主流方案但在端側(cè)推理、模型參數(shù)固定的應(yīng)用場(chǎng)景比如離線(xiàn)推薦、語(yǔ)音喚醒里有機(jī)會(huì)率先落地因?yàn)樗鼘?duì)精度要求相對(duì)寬松、對(duì)功耗極度敏感、且模型一般不常變。6.2 chiplet與異構(gòu)集成把不同工藝的芯片拼成一個(gè)系統(tǒng)傳統(tǒng)的單芯片方案Monolithic在所有功能上都使用同一種工藝但AI芯片里不同模塊對(duì)工藝的要求差異很大——邏輯運(yùn)算需要先進(jìn)工藝的高頻率SRAM需要高密度DRAM的工藝則完全不同。Chiplet思路的核心就是“不再?gòu)?qiáng)求所有功能做在同一塊die上”而是把CPU核、NPU陣列、IO、SRAM分別做成小芯片再用先進(jìn)封裝比如CoWoS、InFO-LSI把它們高密度地拼裝到一個(gè)基板上。這樣做的好處有兩個(gè)一是良率改善明顯小芯片的制造良率遠(yuǎn)高于一顆巨大的SoC所以綜合成本反而更低二是設(shè)計(jì)靈活性提升可以像搭積木一樣組合不同能力的芯片模塊。但配套的難點(diǎn)也在這里——芯片之間的die-to-die接口協(xié)議比如UCIe必須足夠成熟信號(hào)完整性、電源完整性、散熱在封裝級(jí)的協(xié)同設(shè)計(jì)非常復(fù)雜。chiplet本身倒是早就大規(guī)模商用了AMD的EPYC服務(wù)器處理器就是多個(gè)chiplet拼出來(lái)的只是AI芯片場(chǎng)景下的chiplet互連對(duì)帶寬的訴求更加激進(jìn)。6.3 軟件定義硬件可重構(gòu)計(jì)算架構(gòu)是否卷土重來(lái)FPGA是“可重構(gòu)”的元老但傳統(tǒng)FPGA因?yàn)樵贏I計(jì)算上的面積效率太差一直沒(méi)能在主流AI市場(chǎng)打開(kāi)局面。但“軟件定義硬件”這個(gè)理念本身沒(méi)有死反而在新興的架構(gòu)形態(tài)里重新冒頭。一個(gè)是動(dòng)態(tài)可重構(gòu)的NPU陣列芯片的互聯(lián)網(wǎng)絡(luò)和計(jì)算單元的配置可以在運(yùn)行周期內(nèi)修改讓硬件拓?fù)鋭?dòng)態(tài)適配不同模型的算子依賴(lài)關(guān)系。這本質(zhì)上是對(duì)VLIW哲學(xué)的一種延伸——把“硬件怎么連”的決策權(quán)交給編譯器而編譯器可以根據(jù)當(dāng)前模型結(jié)構(gòu)實(shí)時(shí)生成配置位流。另一個(gè)是更徹底的“空域計(jì)算”Spatial Computing方案典型代表是SambaNova家的架構(gòu)。它把計(jì)算圖直接映射到硬件資源的三維結(jié)構(gòu)中數(shù)據(jù)流在計(jì)算單元間流動(dòng)而不是被“取指-譯碼-執(zhí)行”的指令流水線(xiàn)驅(qū)動(dòng)。這種架構(gòu)在編譯器和運(yùn)行時(shí)上投入非常大但一旦映射成功計(jì)算效率確實(shí)能做到非常接近理論上限。這類(lèi)方案的軟件棧復(fù)雜度極高短時(shí)間很難成為主流但它代表了一種重要的思路通用性和高性能并不是天然矛盾的只要你愿意讓軟件做更多事。從架構(gòu)演進(jìn)的角度看未來(lái)三到五年的AI芯片大概率不會(huì)是某一種“終極架構(gòu)”勝出而是多種架構(gòu)在不同場(chǎng)景里長(zhǎng)期共存。云端訓(xùn)練繼續(xù)由GPU和GPU-like方案主導(dǎo)端側(cè)推理被更多專(zhuān)用NPU和存內(nèi)計(jì)算蠶食數(shù)據(jù)中心邊緣則可能出現(xiàn)更多為特定模型定制的“半通用”芯片。7. 一句話(huà)看懂一款A(yù)I芯片架構(gòu)的實(shí)用框架最后分享一個(gè)我自己常用的框架幫助技術(shù)決策者快速判斷一款A(yù)I芯片架構(gòu)的成色。不需要讀懂每一份Datasheet抓住下面五個(gè)維度就能避免大部分選型陷阱。第一是看片上SRAM的容量和帶寬。這決定了單算子能裝多大的tile、算子融合能做到什么程度。低于一定閾值的芯片跑Transformer類(lèi)模型會(huì)非常吃力。第二是看MAC陣列的數(shù)據(jù)流形態(tài)。是脈動(dòng)陣列還是SIMD風(fēng)格還是可重構(gòu)陣列每種形態(tài)都有配套的編譯難度要結(jié)合團(tuán)隊(duì)自己的軟件能力選擇。如果團(tuán)隊(duì)主要靠C/CUDA生態(tài)SIMD風(fēng)格上手更快如果團(tuán)隊(duì)能接受更接近硬件編程模型的學(xué)習(xí)成本脈動(dòng)陣列的上限更高。第三是看軟件棧的支持深度。一個(gè)成熟的AI芯片不只是硬件好編譯器、算子庫(kù)、通信庫(kù)、調(diào)試工具這些“軟件底盤(pán)”決定了一個(gè)小團(tuán)隊(duì)能不能真正用起來(lái)。我經(jīng)常用“從拿到開(kāi)發(fā)板到跑通第一個(gè)Transformer推理需要多久”來(lái)衡量這個(gè)芯片軟件棧的成熟度。第四是看集群互連能力。芯片之間是標(biāo)準(zhǔn)以太網(wǎng)、RoCE還是專(zhuān)用高速互連是否原生支持分布式訓(xùn)練框架的通信原語(yǔ)這決定了你從單卡擴(kuò)展到多卡的平滑程度。第五是看生態(tài)綁定程度。芯片背后的生態(tài)是封閉的還是開(kāi)放的底層編程接口是不是對(duì)上層算法足夠友好這些看似“軟性”的指標(biāo)在實(shí)際生產(chǎn)效率上的影響往往比某些硬指標(biāo)更關(guān)鍵。架構(gòu)設(shè)計(jì)從來(lái)都是在“通用性”和“效率”之間走鋼絲。芯片設(shè)計(jì)如此軟件工程如此做一個(gè)好的技術(shù)選型決策也是如此。評(píng)估一款A(yù)I芯片架構(gòu)不只是在評(píng)估一塊硅片而是在評(píng)估一個(gè)團(tuán)隊(duì)對(duì)計(jì)算本質(zhì)的理解深度以及那個(gè)團(tuán)隊(duì)選擇把未來(lái)押在哪種計(jì)算哲學(xué)上。