AI底層執(zhí)行邏輯:張量與NPU的深度解析)
端側(cè)AI這個詞這兩年出現(xiàn)的頻率越來越高但很多人對它的理解還停留在把模型塞進(jìn)手機(jī)里跑這個層面。真正做過端側(cè)部署的人都知道事情遠(yuǎn)沒有這么簡單。一個模型從訓(xùn)練框架里導(dǎo)出到最終在設(shè)備上以可接受的延遲和功耗跑起來中間要經(jīng)過圖優(yōu)化、算子融合、量化、內(nèi)存布局調(diào)整、硬件指令映射等一長串環(huán)節(jié)。而串聯(lián)這一切的核心概念就是張量和NPU。這篇文章不打算泛泛而談端側(cè)AI的前景而是想把底層執(zhí)行邏輯這條鏈路拆開來看。張量到底是怎么在硬件上流動的NPU和CPU、GPU在處理張量時有什么本質(zhì)區(qū)別為什么同一個模型在不同芯片上性能差距能有好幾倍這些問題的答案藏在數(shù)據(jù)布局、指令調(diào)度和內(nèi)存層級這些不那么性感的細(xì)節(jié)里。如果你正在做端側(cè)部署或者準(zhǔn)備把模型搬到邊緣設(shè)備上這些底層邏輯搞清楚了很多性能問題就能自己定位而不是靠反復(fù)試參數(shù)碰運(yùn)氣。1. 張量不是多維數(shù)組這么簡單1.1 從數(shù)學(xué)對象到內(nèi)存布局的轉(zhuǎn)換大多數(shù)教程在介紹張量時會告訴你張量就是多維數(shù)組。這個說法在PyTorch里寫代碼時沒毛病但一旦進(jìn)入端側(cè)部署環(huán)節(jié)這個理解就會害死人。張量在數(shù)學(xué)上是一個帶有坐標(biāo)變換規(guī)則的多維數(shù)組但在硬件層面它首先是一塊連續(xù)或非連續(xù)的內(nèi)存區(qū)域外加一組描述這塊內(nèi)存如何解釋的元數(shù)據(jù)。舉個具體的例子。一個形狀為[1, 3, 224, 224]的浮點(diǎn)張量在NCHW布局下內(nèi)存里是這樣排列的先存第一個通道的所有像素再存第二個通道以此類推。但在NHWC布局下內(nèi)存里先存第一個像素的三個通道值再存第二個像素的三個通道值。同樣是那個張量數(shù)學(xué)含義完全一樣但內(nèi)存訪問模式截然不同。為什么這件事在端側(cè)這么重要因?yàn)镹PU的矩陣計(jì)算單元通常對內(nèi)存訪問模式有強(qiáng)偏好。很多NPU的卷積加速器在設(shè)計(jì)時就假設(shè)輸入是NHWC布局因?yàn)檫@樣每次讀取連續(xù)內(nèi)存就能拿到一個像素的所有通道值方便做向量化加載。如果你傳進(jìn)去的是NCHW硬件要么需要額外的轉(zhuǎn)置操作要么直接退化成低效的逐元素訪問。我在實(shí)際項(xiàng)目里遇到過一個典型案例同一個MobileNetV2模型在某個NPU上NCHW布局推理耗時38毫秒換成NHWC后降到21毫秒。模型沒變量化參數(shù)沒變僅僅是數(shù)據(jù)布局調(diào)整性能差了將近一倍。這就是張量不只是多維數(shù)組的現(xiàn)實(shí)含義。1.2 步長、偏移與視圖張量的隱藏成本張量的元數(shù)據(jù)里除了形狀shape還有兩個關(guān)鍵字段步長stride和偏移offset。步長描述的是沿著某個維度移動一個位置內(nèi)存地址需要跳多少個元素。偏移描述的是張量數(shù)據(jù)相對于底層內(nèi)存塊的起始位置。這兩個字段的存在使得張量可以是一個視圖而不是一塊獨(dú)立內(nèi)存。比如對一個[1, 3, 224, 224]的張量做轉(zhuǎn)置得到[1, 224, 224, 3]在PyTorch里這個操作幾乎是零成本的因?yàn)樗皇歉牧瞬介L元數(shù)據(jù)底層內(nèi)存沒動。但問題來了這個轉(zhuǎn)置后的張量在內(nèi)存里并不是連續(xù)的當(dāng)你把它傳給NPU時硬件可能無法直接處理非連續(xù)內(nèi)存。端側(cè)部署框架通常會在圖優(yōu)化階段插入連續(xù)化操作把非連續(xù)張量復(fù)制成連續(xù)內(nèi)存。這個復(fù)制操作本身有開銷而且會額外占用內(nèi)存。在內(nèi)存緊張的嵌入式設(shè)備上這種隱藏成本可能直接導(dǎo)致OOM。提示在導(dǎo)出模型前盡量確保所有張量都是連續(xù)的。可以用tensor.is_contiguous()檢查必要時調(diào)用.contiguous()。但要注意過度調(diào)用.contiguous()會引入不必要的復(fù)制最好在圖優(yōu)化階段統(tǒng)一處理。1.3 量化張量的特殊之處端側(cè)AI繞不開量化。一個FP32張量占4字節(jié)每元素量化到INT8后只占1字節(jié)內(nèi)存占用直接降到四分之一而且整數(shù)運(yùn)算在大多數(shù)NPU上比浮點(diǎn)運(yùn)算快得多。但量化張量的內(nèi)存布局和浮點(diǎn)張量有本質(zhì)區(qū)別。量化張量通常采用每通道量化per-channel quantization或每張量量化per-tensor quantization。每通道量化意味著每個通道有獨(dú)立的縮放因子scale和零點(diǎn)zero point。這些參數(shù)需要和量化后的整數(shù)值一起存儲硬件在計(jì)算時需要先讀取scale和zero point再做反量化或直接做整數(shù)運(yùn)算。這里有個容易踩的坑不同框架對量化參數(shù)的存儲順序和內(nèi)存對齊要求不一樣。比如TFLite的量化張量要求scale和zero point按通道連續(xù)存儲而某些NPU工具鏈要求它們按特定對齊方式存放。如果你手動構(gòu)造量化張量對齊沒做好硬件可能讀到錯誤的值導(dǎo)致推理結(jié)果完全不對但又不報錯。這種問題排查起來非常痛苦因?yàn)槟P湍芘芡ㄖ皇墙Y(jié)果錯了。2. NPU到底在算什么2.1 從CPU的標(biāo)量思維切換到NPU的矩陣思維CPU的設(shè)計(jì)哲學(xué)是低延遲、通用性。它有復(fù)雜的控制邏輯、多級緩存、亂序執(zhí)行、分支預(yù)測擅長處理邏輯復(fù)雜的標(biāo)量運(yùn)算。但當(dāng)你讓CPU去做矩陣乘法時它其實(shí)是在用標(biāo)量運(yùn)算模擬矩陣運(yùn)算效率很低。NPU的設(shè)計(jì)哲學(xué)完全不同高吞吐、專用性。它通常包含大量的乘加單元MAC這些單元以脈動陣列systolic array或類似結(jié)構(gòu)組織能夠在每個時鐘周期完成成百上千次乘加運(yùn)算。但代價是靈活性差只擅長處理特定形狀和類型的張量運(yùn)算。舉個例子。一個[1, 256, 256, 64]的卷積層卷積核是[3, 3, 64, 128]。在CPU上這需要嵌套循環(huán)逐元素計(jì)算即使有SIMD指令加速也很難跑滿。在NPU上這個卷積會被映射到矩陣乘法單元輸入張量和權(quán)重張量被分塊加載到片上緩存MAC陣列并行計(jì)算理論上可以在幾十個周期內(nèi)完成。但NPU的矩陣單元通常有固定的尺寸比如128x128或256x256。如果你的張量形狀不匹配就需要padding或分塊。分塊會引入額外的數(shù)據(jù)搬運(yùn)和邊界處理開銷。這就是為什么有些模型在NPU上跑得飛快有些卻還不如CPU——形狀不匹配導(dǎo)致硬件利用率極低。2.2 片上內(nèi)存NPU性能的真正瓶頸很多人以為NPU的性能瓶頸在算力其實(shí)大多數(shù)時候瓶頸在內(nèi)存帶寬和片上緩存容量。NPU的MAC陣列運(yùn)算速度極快但數(shù)據(jù)從DRAM搬到片上緩存的速度遠(yuǎn)遠(yuǎn)跟不上。如果每次計(jì)算都需要從DRAM讀數(shù)據(jù)MAC陣列大部分時間都在等數(shù)據(jù)利用率可能只有百分之十幾。這就是所謂的內(nèi)存墻問題。為了解決這個問題NPU通常采用**分塊tiling**策略把大張量切成小塊每塊能放進(jìn)片上緩存然后在緩存內(nèi)完成計(jì)算減少DRAM訪問次數(shù)。分塊的大小和順序直接影響性能。我實(shí)測過一個典型的端側(cè)模型在某個NPU上默認(rèn)分塊策略下推理耗時45毫秒手動調(diào)整分塊參數(shù)后降到28毫秒。調(diào)整的內(nèi)容包括增大輸入分塊的行數(shù)、調(diào)整權(quán)重分塊的通道數(shù)、改變分塊的計(jì)算順序以減少緩存換入換出。這些參數(shù)在工具鏈里通常有默認(rèn)值但默認(rèn)值不一定適合你的模型。分塊策略片上緩存命中率推理耗時功耗默認(rèn)分塊62%45ms1.8W增大輸入分塊78%34ms1.5W調(diào)整權(quán)重分塊85%28ms1.3W綜合優(yōu)化91%24ms1.2W這張表的數(shù)據(jù)來自我在一個邊緣盒子上的實(shí)測模型是量化后的YOLOv5s。可以看到分塊策略優(yōu)化后不僅速度提升近一倍功耗也明顯下降因?yàn)闇p少了DRAM訪問次數(shù)。2.3 指令調(diào)度與流水線并行NPU內(nèi)部通常有多個計(jì)算單元卷積單元、池化單元、激活單元、量化單元等。這些單元可以流水線并行工作。比如卷積單元在計(jì)算當(dāng)前層時激活單元可以處理上一層的輸出量化單元可以準(zhǔn)備下一層的輸入。但流水線并行需要編譯器或運(yùn)行時做精細(xì)的指令調(diào)度。如果調(diào)度不好單元之間會互相等待流水線出現(xiàn)氣泡性能下降。端側(cè)部署工具鏈通常會自動做指令調(diào)度但調(diào)度質(zhì)量參差不齊。有些工具鏈的調(diào)度器比較保守為了保證正確性插入了過多的同步操作導(dǎo)致并行度不夠。這時候可能需要手動干預(yù)比如調(diào)整層的執(zhí)行順序、合并某些操作、或者用工具鏈提供的調(diào)度提示scheduling hint來引導(dǎo)編譯器。注意手動調(diào)整指令調(diào)度有風(fēng)險可能導(dǎo)致結(jié)果不正確。建議在調(diào)整后做完整的精度驗(yàn)證不要只看推理速度。3. 從框架到硬件的完整鏈路3.1 模型導(dǎo)出那些容易丟掉的元信息從PyTorch或TensorFlow導(dǎo)出模型時很多元信息會丟失。比如PyTorch的torch.jit.trace只記錄張量的形狀和數(shù)據(jù)類型不記錄控制流torch.jit.script能保留控制流但導(dǎo)出的圖可能包含硬件不支持的操作。更隱蔽的是張量命名和布局信息的丟失。在訓(xùn)練框架里張量的維度順序是有明確語義的比如NCHW。但導(dǎo)出到ONNX后如果沒顯式指定布局某些轉(zhuǎn)換工具會默認(rèn)按NHWC處理導(dǎo)致維度錯位。這種錯誤在模型能跑通的情況下很難發(fā)現(xiàn)因?yàn)樾螤羁赡芘銮蓪Φ蒙系?jì)算結(jié)果完全錯誤。我的經(jīng)驗(yàn)是導(dǎo)出模型后一定要用工具如Netron可視化檢查每個節(jié)點(diǎn)的輸入輸出形狀和數(shù)據(jù)類型和原始模型逐一對比。特別是第一個卷積層和最后一個全連接層這兩層最容易出問題。3.2 圖優(yōu)化算子融合的收益與代價圖優(yōu)化階段最常見的操作是算子融合operator fusion。比如把Conv BatchNorm ReLU融合成一個算子。融合的好處是減少中間張量的內(nèi)存讀寫降低kernel啟動開銷。但融合也有代價。融合后的算子對硬件的要求更高如果NPU不支持融合后的算子工具鏈可能回退到CPU執(zhí)行反而更慢。另外融合會改變數(shù)值計(jì)算的順序可能引入微小的精度差異。對于量化模型這種差異可能被放大。我在一個項(xiàng)目里遇到過這樣的情況Conv BN ReLU融合后在某個NPU上精度下降了2個百分點(diǎn)。排查后發(fā)現(xiàn)是BN的縮放因子在融合時被量化到INT8精度損失累積導(dǎo)致最終結(jié)果偏移。解決方案是保留BN不融合或者用更高精度的量化方案處理BN參數(shù)。3.3 內(nèi)存分配靜態(tài)與動態(tài)的權(quán)衡端側(cè)部署通常采用靜態(tài)內(nèi)存分配即在推理前就確定好所有張量的內(nèi)存地址和大小推理過程中不再動態(tài)分配。這樣做的好處是避免內(nèi)存碎片和分配開銷但要求模型的所有張量形狀在編譯時已知。如果模型有動態(tài)形狀比如輸入分辨率可變靜態(tài)內(nèi)存分配就做不了只能用動態(tài)分配。動態(tài)分配在端側(cè)設(shè)備上開銷很大而且容易導(dǎo)致內(nèi)存碎片。折中方案是多靜態(tài)形狀為幾種常見的輸入形狀分別編譯模型運(yùn)行時根據(jù)實(shí)際輸入選擇對應(yīng)的編譯版本。內(nèi)存分配的另一個問題是內(nèi)存復(fù)用。多個張量如果生命周期不重疊可以復(fù)用同一塊內(nèi)存。工具鏈通常會自動做內(nèi)存復(fù)用分析但分析質(zhì)量取決于圖的復(fù)雜度。對于復(fù)雜的模型手動指定內(nèi)存復(fù)用策略可能比自動分析更高效。4. 端側(cè)部署中的典型性能陷阱4.1 形狀不匹配導(dǎo)致的硬件利用率暴跌NPU的矩陣單元有固定尺寸比如128x128。如果卷積的輸入通道數(shù)是64硬件只能利用一半的MAC單元。如果輸入通道數(shù)是48利用率更低。這種浪費(fèi)在深層網(wǎng)絡(luò)中累積起來非常可觀。解決方案通常有兩個一是通道填充channel padding把通道數(shù)補(bǔ)齊到硬件偏好的倍數(shù)二是通道重排channel shuffle把多個小通道的卷積合并成一個大通道的卷積。通道填充會引入額外的計(jì)算量但硬件利用率提升帶來的收益通常更大。我做過一個對比實(shí)驗(yàn)一個通道數(shù)為48的卷積層在128x128 MAC陣列的NPU上直接推理耗時12毫秒通道填充到64后耗時9毫秒填充到128后耗時11毫秒。填充到64是最優(yōu)的因?yàn)榧忍嵘死寐视譀]有引入太多額外計(jì)算。4.2 量化校準(zhǔn)集的選取偏差量化校準(zhǔn)集的選取對最終精度影響極大。很多人隨便找?guī)装購垐D片做校準(zhǔn)結(jié)果量化后精度掉得厲害。校準(zhǔn)集應(yīng)該盡可能覆蓋實(shí)際部署場景中的數(shù)據(jù)分布。比如做人臉檢測模型校準(zhǔn)集里如果全是正臉部署時遇到側(cè)臉就會出問題。校準(zhǔn)集里應(yīng)該包含各種角度、光照、遮擋情況的人臉。另外校準(zhǔn)集的數(shù)量也有講究太少導(dǎo)致量化參數(shù)估計(jì)不準(zhǔn)太多則校準(zhǔn)時間過長。通常幾百到幾千張是比較合理的范圍。還有一個容易被忽略的點(diǎn)校準(zhǔn)集的預(yù)處理必須和推理時的預(yù)處理完全一致。如果校準(zhǔn)集做了歸一化而推理時沒做量化參數(shù)就會完全錯誤。4.3 多核NPU的任務(wù)劃分高端端側(cè)芯片通常有多個NPU核心。如何把模型劃分到多個核心上并行執(zhí)行是一個復(fù)雜的問題。簡單的按層劃分可能導(dǎo)致核心間通信開銷過大因?yàn)閷优c層之間的張量需要跨核心傳輸。更好的策略是按子圖劃分把關(guān)聯(lián)緊密的層放在同一個核心上減少跨核心數(shù)據(jù)傳輸。但子圖劃分需要編譯器做全局分析不是所有工具鏈都支持。我在一個多核NPU項(xiàng)目里的經(jīng)驗(yàn)是如果工具鏈不支持自動子圖劃分可以手動把模型切成幾個獨(dú)立的子模型每個子模型在一個核心上運(yùn)行子模型之間的接口張量盡量小。這樣雖然損失了一些并行度但避免了跨核心通信的瓶頸。5. 調(diào)試端側(cè)AI問題的實(shí)用手段5.1 逐層對比定位精度問題的利器當(dāng)端側(cè)推理結(jié)果和訓(xùn)練框架不一致時最有效的排查方法是逐層對比。把訓(xùn)練框架每一層的輸出保存下來和端側(cè)推理每一層的輸出做對比找到第一個出現(xiàn)明顯差異的層。具體操作在訓(xùn)練框架里注冊hook保存每層輸出在端側(cè)推理時通過調(diào)試接口導(dǎo)出每層輸出。然后計(jì)算每層的余弦相似度或最大絕對誤差。第一個誤差超過閾值的層就是問題所在。這個方法聽起來簡單但實(shí)際操作中有幾個坑。一是端側(cè)推理的中間張量通常不暴露需要工具鏈支持調(diào)試模式。二是量化模型的中間張量是INT8和FP32對比時需要先反量化。三是某些融合算子沒有對應(yīng)的單層輸出需要臨時關(guān)閉融合。5.2 性能剖析找到真正的瓶頸端側(cè)性能剖析比服務(wù)器端困難得多因?yàn)楹芏喽藗?cè)設(shè)備沒有完善的profiling工具。但基本的思路是一樣的測量每個算子的耗時找到耗時最長的算子。如果工具鏈不支持逐算子profiling可以用二分法把模型從中間切成兩半分別測量耗時找到耗時異常的那一半再繼續(xù)切分。雖然粗糙但能快速定位問題區(qū)域。另一個實(shí)用技巧是替換法把可疑的算子替換成已知高效的算子比如把普通卷積替換成深度可分離卷積看性能是否提升。如果提升明顯說明原算子確實(shí)有問題。5.3 內(nèi)存峰值監(jiān)控避免OOM端側(cè)設(shè)備內(nèi)存有限OOM是常見問題。監(jiān)控內(nèi)存峰值的方法因平臺而異但基本思路是在推理前后記錄內(nèi)存使用量在推理過程中定期采樣。如果內(nèi)存峰值超過預(yù)期通常是因?yàn)橹虚g張量太多或太大。解決方案包括啟用內(nèi)存復(fù)用、減小batch size、降低輸入分辨率、使用更激進(jìn)的量化方案。提示有些NPU工具鏈會在編譯時報告內(nèi)存使用估算但這個估算通常偏樂觀。實(shí)際運(yùn)行時內(nèi)存峰值可能更高因?yàn)榇嬖谂R時緩沖區(qū)和對齊填充。建議預(yù)留20%到30%的內(nèi)存余量。6. 一些不那么標(biāo)準(zhǔn)但很實(shí)用的經(jīng)驗(yàn)6.1 不要迷信工具鏈的默認(rèn)配置大多數(shù)端側(cè)部署工具鏈的默認(rèn)配置是為了能跑通而不是跑得快。默認(rèn)配置通常比較保守比如使用較小的分塊、保留所有中間張量、不做激進(jìn)的算子融合。如果你對性能有要求一定要深入工具鏈的配置項(xiàng)逐個調(diào)整。我習(xí)慣的做法是先用默認(rèn)配置跑通記錄基線性能然后逐項(xiàng)調(diào)整配置每次只改一個參數(shù)觀察性能變化最后組合最優(yōu)參數(shù)。這個過程可能耗時但通常能帶來30%到50%的性能提升。6.2 模型結(jié)構(gòu)設(shè)計(jì)要迎合硬件如果模型是你自己設(shè)計(jì)的在結(jié)構(gòu)設(shè)計(jì)階段就應(yīng)該考慮目標(biāo)硬件的特性。比如目標(biāo)NPU偏好3x3卷積就少用5x5和7x7目標(biāo)NPU的通道數(shù)是8的倍數(shù)就把所有層的通道數(shù)設(shè)為8的倍數(shù)目標(biāo)NPU對深度可分離卷積有專門加速就多用深度可分離卷積。這種硬件感知的模型設(shè)計(jì)比事后優(yōu)化有效得多。當(dāng)然前提是你知道目標(biāo)硬件的特性。獲取這些信息的途徑包括芯片廠商的文檔、工具鏈的優(yōu)化指南、以及自己做的微基準(zhǔn)測試。6.3 版本鎖定避免工具鏈升級帶來的意外端側(cè)部署工具鏈的版本兼容性通常很差。今天能跑通的模型升級工具鏈后可能就跑不通了或者性能大幅下降。我的建議是一旦找到能穩(wěn)定工作的工具鏈版本就鎖定這個版本不要輕易升級。如果必須升級一定要在升級前備份當(dāng)前的工作配置升級后做完整的回歸測試包括精度測試和性能測試。不要只看能跑通就認(rèn)為沒問題。6.4 溫度對性能的影響端側(cè)設(shè)備通常散熱條件有限長時間推理會導(dǎo)致芯片溫度升高觸發(fā)降頻。降頻后性能可能下降30%以上。如果你的應(yīng)用需要持續(xù)推理一定要考慮散熱設(shè)計(jì)或者在軟件層面做溫度監(jiān)控和動態(tài)調(diào)頻。我在一個邊緣計(jì)算項(xiàng)目里遇到過這樣的情況冷啟動時推理耗時25毫秒連續(xù)跑10分鐘后降到35毫秒。后來加了散熱片并調(diào)整了推理任務(wù)的調(diào)度策略避免連續(xù)滿負(fù)荷運(yùn)行性能才穩(wěn)定下來。7. 端側(cè)AI底層執(zhí)行邏輯的未來走向從目前的技術(shù)趨勢看端側(cè)AI的底層執(zhí)行邏輯正在向幾個方向演進(jìn)。一是編譯器和硬件的協(xié)同設(shè)計(jì)越來越緊密工具鏈不再只是把模型映射到硬件而是會根據(jù)硬件特性反向指導(dǎo)模型結(jié)構(gòu)設(shè)計(jì)。二是動態(tài)形狀支持越來越完善靜態(tài)編譯的限制正在被逐步打破。三是多模態(tài)融合對底層執(zhí)行提出了新要求視覺、語音、文本的張量需要在同一套硬件上高效流轉(zhuǎn)。但無論技術(shù)怎么演進(jìn)理解張量在硬件上的流動方式、理解NPU的計(jì)算模式和瓶頸、理解從框架到硬件的完整鏈路這些底層邏輯是不會過時的。工具鏈會變芯片會變但這些基本原理是相通的。把底層邏輯搞清楚了面對新工具、新硬件時你就能快速上手而不是從零開始摸索。我在實(shí)際項(xiàng)目中最大的體會是端側(cè)AI的性能優(yōu)化80%的收益來自對底層執(zhí)行邏輯的理解20%來自具體的調(diào)參技巧。很多人花大量時間試各種參數(shù)組合卻不愿意花時間搞清楚張量是怎么在硬件上流動的。結(jié)果就是換個模型或換個硬件之前試出來的參數(shù)全都沒用了。而理解了底層邏輯的人能夠快速分析新場景下的瓶頸在哪里有針對性地做優(yōu)化。這個能力才是端側(cè)AI工程師真正的核心競爭力。