第一篇我們剖開一個 attention block,看到 49 個基本運算:FULLY_CONNECTED、BATCH_MATMUL、SOFTMAX,還有 SIN、COS、SELECT_V2、SLICE、TRANSPOSE。
現在來面對那個尷尬的問題。你有一顆 NPU,它有卷積和矩陣乘法的 kernel,因為那是任何模型裡 90% 的 FLOPs,晶片就是繞著這個設計的。那它有 SIN 的 kernel 嗎?
也許有,也許沒有。如果沒有,這個模型會怎麼樣?總不會直接不能跑吧。
切圖(Partitioning)
圖會被切成幾塊,每塊丟給跑得動它的東西:
[FC][FC][FC] [SIN][COS] [MATMUL][SOFTMAX][MATMUL][FC]
└─ subgraph A ─┘└─ B ──┘ └────── subgraph C ─────────┘
NPU CPU NPU
這就是 graph partitioning,把「這裡有一個模型、那裡有一些硬體」變成一份可執行計畫的那個步驟。
關於它有兩件事值得記住,兩件都不直覺。
每一刀都有代價
直覺的想法是:切圖只是免費的分類工作,把 op 丟進不同的桶子,各自在對的裝置上跑,就這樣。我一開始也是這樣想的。但它不免費,因為這些塊之間根本不共用記憶體。
subgraph A 的輸出住在 NPU 的記憶體裡,用的是 NPU 想要的排列方式,CPU 不見得能直接讀。所以在 SIN 能執行之前,可能需要一次拷貝、一次 layout 轉換,或兩者都要,算完之後還得再送回去。
也就是說,同一張圖切成 2 塊和切成 20 塊,行為完全不同。極端情況下,碎片化嚴重的圖反而會比純 CPU 還慢。你付光了所有搬運成本,卻從來沒在加速器上累積到足夠的工作量把它賺回來。
做過加速管線的人都認得這個死法:NPU 很忙、profiler 看起來也還好,但 end-to-end 的數字比你原本想打敗的 CPU baseline 更差。你會開始懷疑是不是哪裡設定錯了,其實沒有,就是切太碎。
這也是為什麼 LiteRT runtime 裡的 buffer 協商不只是接管線的雜活。Dispatch 層要求每個加速器宣告自己接受哪些 buffer 型別,而文件明講:兩個加速器沒有交集時,runtime 會插入一次轉換。出自 DISPATCH_API.md:
User -> GPU: 你接受哪些 buffer?
GPU -> User: GlTexture, ClBuffer
User -> NPU: 你接受哪些 buffer?
NPU -> User: AHWB
Buffer conversion might happen (GlTexture -> AHWB)
「Zero-copy 管線」是一條路徑的性質,不是某個元件的性質。它只在每一段接受的 buffer 型別真的有交集時才成立。把 GPU 前處理串到一個 buffer 型別不重疊的 NPU 推理,你就悄悄地把一次拷貝放回了每一幀裡。
誰決定,什麼時候決定
第二件不直覺的事:切圖不是晶片廠的工作。
JIT_COMPILATION.md 描述了整個流程,用字很精確。Compiler plugin 編譯的是「the partitioned subgraphs」。切是 framework 做的,vendor plugin 收到的是已經切好的子圖,它只負責為自家硬體編譯。
這個分工是對的。晶片廠不該被要求寫一個圖切割器,他們只需要回答「這個 op 我支不支援」以及「這是我支援的那些 op 編譯出來的碼」。
「什麼時候切和編譯」則是另一個獨立的維度,LiteRT 兩種都支援:
| AOT(提前編譯) | JIT(裝置上編譯) | |
|---|---|---|
| plugin 回傳 | 硬體專屬 bytecode | 一個 opaque JIT handle |
| framework 動作 | 把 bytecode 序列化進 .tflite | 註冊一個空 buffer,什麼都不序列化 |
| dispatch 收到 | kLiteRtDispatchExecutableTypeMlModel | kLiteRtDispatchExecutableTypeJitHandle |
| 適合 | 大模型、目標 SoC 已知 | 小模型、平台無關的發布 |
官方文件列出五家 NPU vendor,而且不是每家兩種都支援。寫作當下 Google Tensor 在 beta 階段是 AOT only,Qualcomm、MediaTek、Intel、Samsung 兩條路都支援。
那個很容易掃過去的 JIT 細節
LiteRT 文件寫得很白,白到我差點直接滑過去,因為它就夾在一段講快取的敘述中間:
Since JIT compilation relies on in-memory handles that exist only during the lifetime of the compiler plugin and the runtime process, JIT compilation cannot be cached.
runtime 一旦偵測到 JIT handle,就會對該次執行關閉 model caching,並記錄 JIT execution handles detected. Disabling JIT model caching.,而且即使你設定了編譯快取目錄,一樣會被繞過。
文件把 JIT 的缺點描述成「較高的首次執行成本(a higher first-run cost)」。這說法準確,但很容易被讀得太寬容。在手機或機上盒上,app 會被啟動、切到背景、在記憶體壓力下被回收,然後再啟動。如果沒有任何東西能跨 process 生命週期活下來,那個「首次執行」成本就是你每次執行都要付的成本。
所以 AOT / JIT 的取捨不是單純的「AOT 需要 toolchain、JIT 不用,所以 JIT 比較方便」:
| AOT | JIT | |
|---|---|---|
| build 時需要 vendor compiler | 需要 | 不需要 |
| 模型可平台無關發布 | 否,每顆 SoC 一份 | 是 |
| 啟動成本 | 低,可快取 | 每次 process 啟動都付,不可快取 |
在一台 process 常被 kill 又重啟的記憶體吃緊裝置上,最後一列可能壓過前面兩列。
這改變了哪些實際做法
三件我開始改變的事。
推理比預期慢的時候,我現在會先問「這張圖被切成幾塊」,再去動別的。碎片化比 kernel 慢常見多了,而且你不主動去看就完全看不到,profiler 只會告訴你每個 op 花多久,不會跳出來說「欸你這裡被切了 17 刀」。runtime 內建 profiler,model-explorer 可以幫你把圖畫出來。
Vendor 的 op 支援清單應該放進選模型的階段,不是評估報告最後的一條附註。知道哪些 op 會造成切斷,你就能挑一個或改一個模型繞開它們,這比事後調參有效太多。回想第一篇,RoPE 是九個基本 op,所以晶片廠怎麼處理這個 pattern,比它標榜的 TOPS 數字更重要。
評估晶片時要問它的 NPU 接受哪些 buffer 型別,不是只問有沒有 NPU。如果那些型別和你的相機或 GPU 前處理路徑沒有交集,你每一幀都會付一次轉換,再高的 TOPS 也救不了。
第三篇再往下一層,進到 Dispatch API:晶片廠要把 NPU 接進 LiteRT 實際上要實作的那份 C 介面、它為什麼取代了 TFLite Delegate,以及那份介面告訴我們在押注一顆晶片之前該問廠商什麼。
資料來源為 google-ai-edge/LiteRT commit dc32e93 的設計文件,JIT_COMPILATION.md、DISPATCH_API.md、COMPILER_PLUGIN.md,以及 Google AI Edge 開發者網站的 NPU acceleration 頁面。