如果只看參數(shù)表你很容易以為手里這顆AI芯片已經(jīng)抵達(dá)性能天花板但當(dāng)你真把一個(gè)復(fù)雜的Transformer或擴(kuò)散模型往上放才會(huì)發(fā)現(xiàn)它的表現(xiàn)和營(yíng)銷材料之間隔著一條深溝。芯片本身沒有錯(cuò)真正的問題是芯片和上層深度學(xué)習(xí)框架之間缺了一層能把模型吃懂、把計(jì)算拆細(xì)、把性能一點(diǎn)點(diǎn)摳出來的軟件引擎。做AI軟件棧這么多年我越來越清楚一件事——這類異構(gòu)計(jì)算架構(gòu)才是決定一顆計(jì)算芯片“核芯”競(jìng)爭(zhēng)力的關(guān)鍵底座。這篇文章不聊某個(gè)具體芯片的跑分而是從一個(gè)實(shí)際做過AI異構(gòu)計(jì)算軟件棧原型的人的角度拆解一顆“核芯引擎”到底由哪幾部分組成、為什么算子融合能帶來肉眼可見的提速、訓(xùn)練到推理轉(zhuǎn)換時(shí)有哪些設(shè)計(jì)取舍以及我真實(shí)踩過的一個(gè)編譯鏈路深坑。如果你正在做算子開發(fā)、推理部署、圖編譯或AI平臺(tái)底層相關(guān)工作這篇東西應(yīng)該能幫你少走幾條彎路。1. 為什么AI硬件必須有一個(gè)專屬引擎而沒人能直接繞過它1.1 從“能用驅(qū)動(dòng)”到“好用的計(jì)算架構(gòu)”先講個(gè)我常拿來舉的例子。很多AI加速卡剛流片回來時(shí)廠商總會(huì)給一套底層的驅(qū)動(dòng)和幾個(gè)示例程序跑個(gè)圖像分類網(wǎng)絡(luò)完全沒有問題。但只要用戶把網(wǎng)絡(luò)換成帶注意力機(jī)制、動(dòng)態(tài)Shape或復(fù)雜控制流的模型性能就會(huì)突然離譜起來有時(shí)甚至比CPU還慢。原因很簡(jiǎn)單AI專用芯片不是通用處理器它的指令集往往非?！罢敝簧瞄L(zhǎng)矩陣乘法、卷積、向量操作等固定模式。它沒有現(xiàn)代CPU那種復(fù)雜的亂序執(zhí)行、分支預(yù)測(cè)和緩存體系也不會(huì)自動(dòng)幫你把復(fù)雜計(jì)算拆成適合并行執(zhí)行的子任務(wù)。芯片能做的只是執(zhí)行一個(gè)個(gè)底層算子而“如何把模型翻譯成底層算子”“怎么調(diào)度這些算子在多核上跑”“中間數(shù)據(jù)放在哪里才不浪費(fèi)帶寬”這些問題全都要靠芯片之上的計(jì)算架構(gòu)引擎去解決。這層架構(gòu)放在整體鏈路里的位置大致是這樣的最上層是AI框架負(fù)責(zé)定義模型結(jié)構(gòu)、反向傳播邏輯和訓(xùn)練流程。中間這層架構(gòu)引擎負(fù)責(zé)接收框架傳來的計(jì)算圖做圖分析、算子映射、圖優(yōu)化、資源編排。最下層是芯片驅(qū)動(dòng)和硬件抽象層負(fù)責(zé)把指令真正發(fā)到計(jì)算單元上。你可以把AI框架理解為菜譜把芯片理解為一口已經(jīng)設(shè)計(jì)得非常精巧的鍋而這層架構(gòu)引擎就是那個(gè)在后廚控制火力、安排上菜順序、調(diào)配廚師時(shí)間的后廚總管。沒有總管鍋再好也只能做一道菜做不了整桌宴席。1.2 一張補(bǔ)全硬件弱點(diǎn)的軟件底版硬件設(shè)計(jì)師在流片時(shí)會(huì)做很多取舍算力做得高內(nèi)存帶寬可能就會(huì)卡脖子通用性做得強(qiáng)單位功耗的性能可能就不好看。芯片一旦成型很多硬件層面的“性格缺陷”就改不了了只能靠上層軟件去補(bǔ)。常見到需要軟件去補(bǔ)的東西我隨手就能列出一堆內(nèi)存層次很淺多級(jí)Cache缺失需要引擎盡可能讓數(shù)據(jù)待在最靠近計(jì)算單元的地方。計(jì)算單元之間的同步機(jī)制原始通常只有事件和信號(hào)量需要引擎來編排異步執(zhí)行順序。不同算子對(duì)數(shù)據(jù)排布格式有不同偏好NCHW和NHWC混用需要引擎在中間做布局轉(zhuǎn)換。多芯片多卡通信時(shí)集合通信的拓?fù)浜颓蟹植呗砸残枰鎭斫y(tǒng)一規(guī)劃。所以你會(huì)發(fā)現(xiàn)真正能支持一個(gè)模型高效跑起來的從來不是一塊芯片而是一整套把芯片能力接住、補(bǔ)齊、放大、交付出去的軟件底版。這塊底版設(shè)計(jì)得好不好才是很多AI項(xiàng)目能不能落地的分水嶺。2. 引擎的真正心臟圖編譯、算子庫(kù)、運(yùn)行時(shí)三者的協(xié)作機(jī)制2.1 圖編譯把模型切成硬件認(rèn)識(shí)的原語(yǔ)圖編譯簡(jiǎn)單說就是把深度學(xué)習(xí)模型的計(jì)算圖從框架側(cè)那種偏邏輯描述的形式逐步降級(jí)成硬件能直接執(zhí)行的原語(yǔ)序列。這個(gè)環(huán)節(jié)通常分好幾層來做不是一步到位的大概流程可以理解為下面這張圖graph load_model(attention_model.onnx) # 第一輪形狀推導(dǎo)和信息補(bǔ)全把每個(gè)中間張量的維度、數(shù)據(jù)類型標(biāo)出來 graph infer_shape(graph) # 第二輪布局轉(zhuǎn)換把模型里不適合硬件訪存的格式調(diào)整為硬件偏好格式 graph layout_convert(graph) # 第三輪圖優(yōu)化尋找可以融合的算子組合減少內(nèi)存搬運(yùn)和內(nèi)核啟動(dòng)開銷 graph fuse_patterns(graph, patterns[conv_bn_relu, transpose_conv]) # 第四輪圖切分將圖分配到不同設(shè)備并插入跨設(shè)備通信原語(yǔ) graph split_for_device(graph, device_count8)很多新手不理解為什么不能直接把模型“一股腦”扔給芯片。核心原因是芯片的執(zhí)行粒度太小模型里一個(gè)Conv層在芯片上要展開成循環(huán)、矩陣分塊、數(shù)據(jù)搬運(yùn)、累加等多個(gè)階段。圖編譯的作用就是把框架層的“大算子”拆成硬件層的“小原語(yǔ)”同時(shí)盡量保留可以合并、可以跳躍的機(jī)會(huì)。我實(shí)際做圖編譯時(shí)最大的體會(huì)是這一層的核心難點(diǎn)不在“翻譯”而在“翻譯之后還能不能找回語(yǔ)義”。一旦你把多個(gè)算子融合成一個(gè)自定義大算子就必須保證反向傳播、斷點(diǎn)調(diào)試、算子性能分析這些工具還能看明白這張圖。否則模型能跑但出了問題誰(shuí)也查不了。2.2 算子庫(kù)與自動(dòng)生成性能的長(zhǎng)期底座圖編譯把圖變成原語(yǔ)序列之后真正干活的還是算子。算子庫(kù)的質(zhì)量直接決定性能上限。早期做AI芯片軟件棧團(tuán)隊(duì)往往會(huì)把大量精力放在人工調(diào)優(yōu)常用算子上卷積、矩陣乘、歸一化、池化、Softmax一個(gè)個(gè)手工寫匯編或內(nèi)聯(lián)匯編。但人工調(diào)優(yōu)有一個(gè)天花板算子太多、Shape組合太多。同一個(gè)卷積可能輸入是四維、五維可能是1x1、3x3、7x7可能是普通卷積、空洞卷積、轉(zhuǎn)置卷積?？咳耸秩ジ采w所有場(chǎng)景根本不現(xiàn)實(shí)。所以現(xiàn)在的主流做法是“人工高優(yōu)算子做標(biāo)桿 自動(dòng)生成工具做兜底”。自動(dòng)生成工具的思路是先定義一個(gè)算子描述語(yǔ)言描述輸入輸出的形狀關(guān)系、計(jì)算邏輯和數(shù)據(jù)布局然后由工具根據(jù)目標(biāo)芯片的寄存器數(shù)量、向量寬度、存儲(chǔ)層級(jí)去搜索最優(yōu)實(shí)現(xiàn)。這個(gè)搜索空間非常大經(jīng)常需要用到代價(jià)模型和啟發(fā)式規(guī)則去剪枝。這中間有一個(gè)我特別想提醒的點(diǎn)自動(dòng)生成出來的算子經(jīng)常在“中等Shape”下表現(xiàn)很好但在極小的Shape下反而會(huì)打不過人工調(diào)優(yōu)版本因?yàn)樗傻拇a前綴開銷太大。我們?cè)谧鏊阕臃职l(fā)時(shí)不能只按算子類型去選擇內(nèi)核還要按實(shí)際Shape、對(duì)齊條件、是否動(dòng)態(tài)Shape去做多層分發(fā)。2.3 運(yùn)行時(shí)調(diào)度把資源和執(zhí)行順序安排明白編譯完成之后運(yùn)行時(shí)負(fù)責(zé)把“計(jì)劃”真正執(zhí)行起來。它管理的東西包括內(nèi)存池、任務(wù)隊(duì)列、事件同步、設(shè)備間的通信編排。很多人會(huì)覺得運(yùn)行時(shí)只是個(gè)“執(zhí)行器”沒什么技術(shù)含量實(shí)際恰恰相反。一個(gè)常見的性能殺手是內(nèi)存碎片。AI模型執(zhí)行過程中要頻繁申請(qǐng)和釋放中間張量如果運(yùn)行時(shí)直接依賴底層系統(tǒng)級(jí)內(nèi)存分配速度會(huì)慢得讓人崩潰。所以一個(gè)合格的引擎會(huì)在啟動(dòng)時(shí)預(yù)先申請(qǐng)一大塊設(shè)備內(nèi)存然后自己做池化和分配策略把常用張量的生命周期盡量重疊減少實(shí)際分配次數(shù)。另一個(gè)坑是異步執(zhí)行。為了隱藏?cái)?shù)據(jù)傳輸延遲引擎會(huì)把計(jì)算和數(shù)據(jù)搬運(yùn)放到不同的執(zhí)行流上。這時(shí)候如果沒有合理的事件依賴關(guān)系就可能出現(xiàn)計(jì)算結(jié)果還沒算完就開始拷貝數(shù)據(jù)的情況。我們內(nèi)部就出過這種問題某些結(jié)構(gòu)跑一萬(wàn)次才偶現(xiàn)一次運(yùn)行錯(cuò)誤排查難度極高最后發(fā)現(xiàn)是事件同步漏了一條支路。最理想的運(yùn)行時(shí)應(yīng)該讓開發(fā)者感覺不到它的存在所有資源管理都在引擎層自動(dòng)完成?？上КF(xiàn)實(shí)是越接近性能極限運(yùn)行時(shí)的存在感就會(huì)越強(qiáng)。3. 性能優(yōu)化最大性價(jià)比算子融合是怎么把速度“摳”出來的3.1 算子融合前必須讀懂的執(zhí)行模型我做了這么久異構(gòu)計(jì)算如果要挑一個(gè)性價(jià)比最高、最常被提及的優(yōu)化手段那就是算子融合。算子融合看著簡(jiǎn)單就是把多個(gè)算子合并成一個(gè)算子但它背后的收益邏輯遠(yuǎn)不止“少調(diào)幾次接口”這么簡(jiǎn)單。先看一次完整算子執(zhí)行的代價(jià)拆解啟動(dòng)開銷內(nèi)核從CPU端提交到設(shè)備端創(chuàng)建任務(wù)、分配上下文通常有幾十微秒級(jí)開銷。讀數(shù)據(jù)從片外內(nèi)存把輸入讀到片上。計(jì)算在計(jì)算單元上做實(shí)際運(yùn)算。寫回把結(jié)果寫回片外內(nèi)存。當(dāng)兩個(gè)算子串行執(zhí)行時(shí)第二個(gè)算子還要再把第一個(gè)算子的輸出從片外讀一遍。以卷積加ReLU為例未融合時(shí)Conv輸出寫回片外內(nèi)存RelU再讀出來白白浪費(fèi)一次帶寬。模型越深這種浪費(fèi)越夸張。未融合Conv - (寫回片外) - ReLU - (寫回片外) 融合后ConvReLU - (只在片內(nèi)搬運(yùn)數(shù)據(jù)最終寫回一次)所以融合的本質(zhì)是把中間結(jié)果留在片上Cache或寄存器里讓數(shù)據(jù)不被“折騰”回片外。這種精細(xì)的偷帶寬動(dòng)作對(duì)AI芯片的加速效果非常明顯尤其帶寬吃緊的推理場(chǎng)景。3.2 常見融合組與格式化對(duì)齊的細(xì)節(jié)融合也不是見算子就合需要滿足形狀相同、依賴關(guān)系簡(jiǎn)單、語(yǔ)義可保持等條件。常見的融合組有這么幾類融合模式說明收益來源Conv BatchNorm ReLU圖像網(wǎng)絡(luò)最經(jīng)典組合少寫回兩次片外減少內(nèi)核啟動(dòng)MatMul Add Activation全連接網(wǎng)絡(luò)常用把偏置和激活融入矩陣乘Transpose Concat / Split處理數(shù)據(jù)搬運(yùn)鏈減少中間布局轉(zhuǎn)換LayerNorm 量化 / 反量化注意力模型后處理減少低比特轉(zhuǎn)換開銷我特別想講的是數(shù)據(jù)格式對(duì)齊。很多模型本來跑得好好的編譯器為了迎合某個(gè)算子的硬件偏好會(huì)突然插入一個(gè) Transpose 或 layout 轉(zhuǎn)換內(nèi)核。這個(gè)轉(zhuǎn)換本身也是算子也要讀寫一遍內(nèi)存。在某些情況下轉(zhuǎn)換引入的開銷比優(yōu)化省下來的還多。我們做圖優(yōu)化時(shí)有一個(gè)原則布局轉(zhuǎn)換不能局部地看必須放到整個(gè)圖里去評(píng)估。一個(gè) Transpose 放在圖的哪個(gè)位置、后面的算子能不能吸收它都必須通過代價(jià)模型來判斷。如果只看單一算子而忽略整體數(shù)據(jù)流很容易陷入局部最優(yōu)。3.3 量化與內(nèi)存對(duì)齊對(duì)編譯決策的影響推理場(chǎng)景進(jìn)來之后低比特量化成了繞不開的話題。INT8算子通常需要輸入張量在某個(gè)對(duì)齊邊界附近還要額外插入量化參數(shù)縮放。這些量化反量化算子天然會(huì)把一段連續(xù)的計(jì)算鏈“切開”導(dǎo)致后面想融合也無從下手。比如一個(gè)常見的注意力結(jié)構(gòu)里Q/K/V做矩陣乘之前要各自做量化Softmax之后又要反量化。如果圖編譯的融合策略只考慮了高精度算子那么這些低比特轉(zhuǎn)換點(diǎn)就會(huì)把圖切成很多碎片造成大量片上片上搬運(yùn)。另一個(gè)被忽略的問題是內(nèi)存對(duì)齊。芯片內(nèi)部有些存儲(chǔ)體是按固定容量組織的比如某些加速器的SRAM一次只能訪問32字節(jié)對(duì)齊的數(shù)據(jù)。如果張量最后的維度不是對(duì)齊的算子內(nèi)核就會(huì)產(chǎn)生大量“壓包”和“解包”操作。這個(gè)問題在動(dòng)態(tài)Shape里尤其嚴(yán)重因?yàn)榫幾g期根本不知道運(yùn)行時(shí)Shape的具體值沒法做精確對(duì)齊只能按最大可能預(yù)留內(nèi)存利用率就會(huì)下降。我的建議是圖編譯器和算子層必須共享同一套對(duì)齊規(guī)則編譯期在插入布局轉(zhuǎn)換時(shí)就要把對(duì)齊信息算清楚而不是把這個(gè)問題拋給算子內(nèi)核去臨時(shí)處理。4. 從訓(xùn)練到推理、從云端到端側(cè)架構(gòu)設(shè)計(jì)必須做好的“三邊適配”4.1 訓(xùn)練、推理對(duì)引擎的不同需求很多人以為引擎只要能跑模型就行訓(xùn)練和推理差別不大。實(shí)際上這兩類場(chǎng)景對(duì)引擎的需求差異非常大。維度訓(xùn)練場(chǎng)景推理場(chǎng)景精度類型FP32、BF16為主容忍范圍嚴(yán)格INT8、FP16為主精度預(yù)算靈活圖規(guī)模大圖、動(dòng)態(tài)Shape、反向邊復(fù)雜圖通常固定越小越好批量大Batch追求吞吐單請(qǐng)求延時(shí)要低Batch常為1錯(cuò)誤容忍出錯(cuò)必須報(bào)錯(cuò)不能跳過偶錯(cuò)可容忍更看重穩(wěn)定我見過最典型的問題是把訓(xùn)練側(cè)的編譯策略直接搬去推理側(cè)用。訓(xùn)練圖有反向邊、有中間梯度張量圖結(jié)構(gòu)復(fù)雜推理圖是前向圖很多中間結(jié)果根本不需要保留。如果引擎不區(qū)分這兩類圖就會(huì)保留大量無用的中間輸出白白浪費(fèi)顯存和帶寬。推理側(cè)還有一個(gè)更“心機(jī)”的優(yōu)化點(diǎn)把Batch設(shè)為1的靜態(tài)圖做極致特化。Batch為1時(shí)很多矩陣乘可以退化成向量化更好的形式內(nèi)存訪問模式也完全不同。做推理部署時(shí)我們經(jīng)常會(huì)根據(jù)具體Batch值重新編譯一份特化圖效果比通用圖要好不少。4.2 動(dòng)態(tài)Shape與靜態(tài)圖的平衡動(dòng)態(tài)Shape對(duì)圖編譯器來說是一個(gè)長(zhǎng)期痛點(diǎn)。做靜態(tài)圖優(yōu)化時(shí)所有張量維度都是確定的編譯器可以提前算好內(nèi)存偏移、規(guī)劃好循環(huán)展開方式。但一旦Shape是動(dòng)態(tài)的很多東西就只能在運(yùn)行時(shí)才能決定圖編譯器不得不插入大量運(yùn)行時(shí)判斷邏輯。這里有個(gè)關(guān)鍵設(shè)計(jì)選擇到底是“一個(gè)動(dòng)態(tài)圖硬跑”還是“根據(jù)常見Shape提前特化幾份版本運(yùn)行時(shí)最接近的復(fù)用”我們后來的做法是采取混合策略引擎先把模型的輸入Shape分成幾種典型模式比如圖像分類常見到224、224注意力模型常見到序列長(zhǎng)度128、256、512。編譯期針對(duì)這些常見Shape生成特化圖并把編譯結(jié)果緩存下來。運(yùn)行時(shí)如果遇到緩存里沒有的新Shape再走一遍動(dòng)態(tài)編譯流程。這種策略在推理服務(wù)里效果很好九成以上的請(qǐng)求都能命中緩存。但這套方案的代價(jià)是代碼復(fù)雜度陡增。緩存命中率、內(nèi)存占用、失效清理這些都是需要細(xì)致處理的工程問題。很多團(tuán)隊(duì)裁掉這些點(diǎn)只在方案PPT里寫“支持動(dòng)態(tài)Shape”上線后才發(fā)現(xiàn)頻繁重新編譯把延遲拉爆了。4.3 端邊云統(tǒng)一部署的取舍一個(gè)架構(gòu)引擎如果能同時(shí)覆蓋云端和邊緣端是很加分的賣點(diǎn)。但云端和端側(cè)的限制條件完全不同云端芯片算力強(qiáng)、內(nèi)存充足、功耗不是首要約束端側(cè)芯片算力弱、內(nèi)存小、功耗敏感還要考慮相對(duì)嚴(yán)苛的散熱條件。為了兼顧兩端引擎通常會(huì)把“編譯目標(biāo)”和“運(yùn)行時(shí)環(huán)境”解耦。云端編譯端側(cè)模型把編譯產(chǎn)物放到端側(cè)執(zhí)行端側(cè)只保留一個(gè)很小的運(yùn)行時(shí)。這個(gè)策略聽起來很美落地時(shí)卻有一堆細(xì)節(jié)端側(cè)內(nèi)存池大小怎么定云端編譯期不能完全預(yù)知端側(cè)的運(yùn)行水位。端側(cè)算子庫(kù)不能覆蓋所有算子編譯時(shí)就要做算子可用性檢查。端側(cè)的批處理、多模型并發(fā)策略和云端完全不一樣統(tǒng)一調(diào)度框架很難同時(shí)討好兩邊。我們?cè)趯?shí)踐中的體會(huì)是不要一開始就追求全場(chǎng)景統(tǒng)一先把一個(gè)場(chǎng)景做透再把核心抽象層抽出來擴(kuò)展。統(tǒng)一架構(gòu)是目標(biāo)但不該是第一步。5. 實(shí)際排障日記一個(gè)算子適配問題暴露出的編譯鏈路短板5.1 問題浮出模型在某個(gè)硬件上偶發(fā)性性能驟降有一次我們?cè)谝粋€(gè)偏重圖像分類的模型集上做了性能回歸結(jié)果很順利大部分模型都能達(dá)到預(yù)期加速比。但后來加入一個(gè)新結(jié)構(gòu)的模型在同樣的芯片上端到端延遲比預(yù)期高了差不多40%。一開始大家都很困惑因?yàn)閱为?dú)看每個(gè)算子的耗時(shí)都挺正常的沒有哪個(gè)算子像“罪魁禍?zhǔn)住?。更討厭的是如果把模型輸入稍微改小一點(diǎn)比如把分辨率從224改成224但這個(gè)頭的維度不同性能問題就消失。典型的偶發(fā)性能事故最讓人頭疼。5.2 定位鏈路從框架到運(yùn)行時(shí)的換段排查我們當(dāng)時(shí)沒有急著改代碼而是先做了一輪層層切分關(guān)閉圖優(yōu)化再跑一遍。如果性能差距縮小說明問題出在圖優(yōu)化/融合階段而不是單個(gè)算子內(nèi)部。打開引擎的圖Dump把編譯前后的計(jì)算圖導(dǎo)出來對(duì)比算子的數(shù)量和類型。用Profiling統(tǒng)計(jì)每個(gè)kernel的執(zhí)行時(shí)間和內(nèi)存搬運(yùn)量重點(diǎn)看有沒有異常的Transpose或Copy操作。對(duì)比Shape推導(dǎo)前后的張量維度變化找出哪個(gè)階段插入的額外轉(zhuǎn)換。最后Profiling結(jié)果指向了一個(gè)很奇怪的點(diǎn)網(wǎng)絡(luò)某一段里出現(xiàn)了一大一小兩個(gè)連續(xù)的Transpose前面那個(gè)還拖著一個(gè)把它輸出寫成“非連續(xù)內(nèi)存”的標(biāo)記。后續(xù)算子為了讓數(shù)據(jù)變得連續(xù)又插入了一次“轉(zhuǎn)連續(xù)”的拷貝操作。這兩個(gè)操作疊加起來在某個(gè)具體的頭部結(jié)構(gòu)里引發(fā)了跨核訪存的沖突延遲直接翻倍。一句話總結(jié)問題根因底層算子匹配器只按固定的數(shù)據(jù)布局模式去匹配卷積算子忽略了這個(gè)模型里特有的布局路徑圖優(yōu)化階段又沒有一個(gè)“布局轉(zhuǎn)換吸收”的規(guī)則去消化它導(dǎo)致引擎不得不在運(yùn)行時(shí)用兩個(gè)額外的復(fù)制算子去兜底。5.3 最終根因與修復(fù)后的復(fù)盤心得修復(fù)其實(shí)不算復(fù)雜我們新增了一條布局轉(zhuǎn)換的優(yōu)化規(guī)則讓編譯器在面對(duì)類似結(jié)構(gòu)時(shí)優(yōu)先選擇可以吸收轉(zhuǎn)換的算子組合并在圖變換階段做全局?jǐn)?shù)據(jù)布局規(guī)劃而不是每次插入Transpose都是局部的即時(shí)行為。這個(gè)坑給我們團(tuán)隊(duì)留下了三個(gè)認(rèn)知層面的教訓(xùn)圖編譯器不能只對(duì)“標(biāo)準(zhǔn)模型”做優(yōu)化必須覆蓋真實(shí)世界里的各種“踩邊”結(jié)構(gòu)。布局轉(zhuǎn)換的代價(jià)必須納入全局代價(jià)模型局部最優(yōu)的插入很可能是全局災(zāi)難。性能回歸測(cè)試不能只看端到端時(shí)間還要定期檢查圖中的Transpose、Copy等“灰色算子”數(shù)量是否膨脹。這種問題在PUBLIC題里很難被看到但恰恰是AI軟件棧工程里最常見的深水區(qū)。編譯優(yōu)化鏈路越長(zhǎng)任何一個(gè)環(huán)節(jié)的疏忽都會(huì)在運(yùn)行時(shí)被放大。6. 判斷一個(gè)計(jì)算架構(gòu)是否“核芯”的經(jīng)驗(yàn)尺度6.1 先看接口抽象再看性能指標(biāo)每次有人讓我評(píng)價(jià)某個(gè)AI計(jì)算架構(gòu)好不好我第一反應(yīng)都不是看它的Benchmark有多漂亮而是看它給開發(fā)者提供的接口抽象是不是清晰、能不能站在開發(fā)者的角度思考問題。一個(gè)合格的引擎應(yīng)該有四層基本接口圖接入接口、算子開發(fā)接口、運(yùn)行時(shí)控制接口、性能分析接口。如果這四個(gè)接口設(shè)計(jì)得清晰開發(fā)者想接入新模型、寫新算子、查性能瓶頸都能很快上手。反之如果接口語(yǔ)義混亂哪怕性能再?gòu)?qiáng)也只是實(shí)驗(yàn)室里的漂亮數(shù)字根本沒法產(chǎn)品化。我特別看重算子開發(fā)接口。因?yàn)樵趯?shí)際項(xiàng)目里總有標(biāo)準(zhǔn)算子庫(kù)覆蓋不了的奇怪結(jié)構(gòu)。如果自定義算子這套體系足夠簡(jiǎn)單團(tuán)隊(duì)自定義算子的成本就會(huì)低很多如果接口復(fù)雜到要懂各種硬件細(xì)節(jié)那優(yōu)先選擇標(biāo)準(zhǔn)算子就會(huì)嚴(yán)重限制模型的自由度。6.2 兼容性與迭代策略計(jì)算架構(gòu)引擎最大的噩夢(mèng)是“今天能用明天就崩”。芯片迭代、AI框架版本升級(jí)、模型結(jié)構(gòu)演進(jìn)任何一個(gè)變化都可能讓之前的優(yōu)化全部失效。一個(gè)穩(wěn)定的引擎必須做到向后兼容和可回退。我們實(shí)踐里的做法是會(huì)為每個(gè)圖優(yōu)化注冊(cè)一個(gè)“規(guī)則ID”每次編譯時(shí)會(huì)把整個(gè)優(yōu)化鏈路版本記錄下來方便定位問題。同時(shí)還要求所有優(yōu)化規(guī)則必須是可禁用的方便在問題出現(xiàn)時(shí)快速縮小范圍。迭代策略上也有一條重要經(jīng)驗(yàn)不要為了趕新功能而頻繁改動(dòng)核心圖變換架構(gòu)。圖變換是引擎的“高速公路”一旦改錯(cuò)所有到這條路線的性能都會(huì)崩盤。更穩(wěn)妥的做法是在旁邊修“輔路”先做小范圍驗(yàn)證再逐步切換主路流量。6.3 對(duì)真實(shí)模型和長(zhǎng)尾場(chǎng)景的支持度一個(gè)引擎值不值得信最終看它對(duì)真實(shí)模型和長(zhǎng)尾場(chǎng)景的照顧程度。很多公開樣例模型都是被調(diào)得干干凈凈的不存在動(dòng)態(tài)控制流、不存在奇怪的Padding、不存在內(nèi)存布局沖突。真實(shí)世界的模型簡(jiǎn)直千奇百怪有的用了多重嵌套控制流有的張量維度是3維或6維有的算子輸出被反復(fù)重新解釋成不同形狀。這些場(chǎng)景才是引擎試金石。我們后來定了一個(gè)例行動(dòng)作每周從線上灰度模型里隨機(jī)抽樣幾十個(gè)放進(jìn)離線回歸環(huán)境跑一遍全流程。這么做不為別的就是防止引擎出現(xiàn)“只認(rèn)識(shí)優(yōu)化過的樣例認(rèn)不出真實(shí)世界”的傲慢。一個(gè)真正有生命力的計(jì)算架構(gòu)引擎不應(yīng)該只對(duì)榜單負(fù)責(zé)。7. 給同樣在做AI軟件棧的同學(xué)的三個(gè)實(shí)操建議7.1 先做窄而深的垂直打通再做橫向抽象剛接觸AI計(jì)算架構(gòu)的人很容易一上來就規(guī)劃一個(gè)大而全的引擎想著框架接口、圖編譯、算子庫(kù)、運(yùn)行時(shí)全都一網(wǎng)打盡。這種項(xiàng)目十有八九會(huì)爛尾因?yàn)槊恳粚拥膹?fù)雜度都超出預(yù)期任何一個(gè)子系統(tǒng)都?jí)蛞粋€(gè)團(tuán)隊(duì)消耗一整年。更務(wù)實(shí)的路線是先選一條極窄的垂直場(chǎng)景比如只做某一個(gè)視覺模型的推理優(yōu)化把這個(gè)模型從接入、圖編譯、算子調(diào)優(yōu)到運(yùn)行時(shí)交付的整條鏈路徹底打通。這條鏈上踩過的每一個(gè)坑都會(huì)讓你對(duì)引擎的設(shè)計(jì)邊界有真正體感。打通之后再做抽象把通用部分抽出來擴(kuò)展比憑空設(shè)計(jì)一套架構(gòu)要可靠得多。7.2 盡早建立Profiling基線別等性能崩了再去查性能問題最怕的不是難查而是沒有基線。如果你不知道某個(gè)算子在硬件上理論能跑多快、不知道數(shù)據(jù)搬運(yùn)在整條鏈路里的延遲分布那么性能一降你連“降在哪”都說不清。我們的做法是做完任何一個(gè)優(yōu)化后立刻記錄三組數(shù)字端到端時(shí)間、算子總耗時(shí)、各階段數(shù)據(jù)搬運(yùn)耗時(shí)。這三組數(shù)字長(zhǎng)期保存下來后續(xù)每次優(yōu)化都跟歷史基線對(duì)比。哪怕優(yōu)化失敗也能快速知道失敗發(fā)生在哪一層。Profiling工具鏈不是優(yōu)化完成后再配的它應(yīng)該是引擎開發(fā)第一天就要有的組成部分。7.3 把回歸用例寫進(jìn)CI用自動(dòng)化封住長(zhǎng)尾問題的口子排障越多越會(huì)發(fā)現(xiàn)大部分性能回退不是“大改進(jìn)”造成的而是某個(gè)小改動(dòng)讓一條長(zhǎng)尾路徑發(fā)生了意料之外的連鎖反應(yīng)。這種問題靠人工Review很難攔住真正有效的手段是把回歸用例硬化成CI門禁。我把融合、布局轉(zhuǎn)換、動(dòng)態(tài)Shape、多設(shè)備切分這些高風(fēng)險(xiǎn)路徑都做成了一組模板測(cè)試每次代碼合并前必須全部跑過并給出性能基線差值的報(bào)告。如果某次改動(dòng)讓某個(gè)關(guān)鍵場(chǎng)景延遲增加超過5%CI就會(huì)直接攔下。這套機(jī)制初期投入不小但它省掉了我們后來幾乎所有“半夜上線事故”的麻煩。說實(shí)話做AI計(jì)算架構(gòu)軟件棧最迷人的地方不是某個(gè)算子能跑多快而是你能站在框架和硬件中間把一層層的抽象打磨得嚴(yán)絲合縫讓模型開發(fā)者感受不到芯片的存在。這種感覺只有真正在那里熬夜調(diào)過Profiling、看過圖Dump、為一兩個(gè)百分點(diǎn)的性能反復(fù)權(quán)衡的人才會(huì)懂。希望這篇把引擎拆開講的筆記也能讓你對(duì)這條“芯”路上的風(fēng)景有個(gè)更清楚的認(rèn)識(shí)。