突破無規則稀疏計算邊界，編譯框架CROSS數倍提升模型性能

AIxiv專欄是機器之心發佈學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯繫報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本篇工作已被 HPCA 2025 接收，由上海交大先進計算機體系結構實驗室蔣力教授課題組（IMPACT）完成，同時也獲得了上海期智研究院的支持。第一作者是劉方鑫老師與博士生黃世遠。

在現代 AI 模型的快速迭代中，如何在保持模型精度的同時提升計算效率成為關鍵課題。尤其在大規模 AI 推理中，非結構化稀疏矩陣的計算效率低下成為難以突破的瓶頸。面對這一挑戰，我們自主研發了 CROSS—— 一種創新的端到端稀疏編譯優化方案，為 AI 推理帶來細粒度稀疏計算的加速效果。

稀疏計算的挑戰：如何處理非均勻稀疏分佈

非結構化細粒度稀疏場景下模型推理效率低下問題是 AI 編譯社區面對的關鍵問題之一。相比於密集算子加速庫（cuBlas），主要的稀疏算子加速庫或編譯框架需要在較高稀疏率下才能獲得收益，而過高的稀疏率需求可能使我們面臨模型精度下降的風險。

圖 1. 相比於 cuBlas，不同稀疏加速庫或編譯框架在不同稀疏率下的加速比。Sputnik、TVM-Sparse、SparseTIR、ASpT 和 cuSPARSE 在稀疏率超過 76%、80.5%、82.6%、89.4% 和 98.1% 時才能獲得正向收益（稀疏矩陣源於 Bert 模型中的稀疏權值矩陣）。

稀疏計算的機會：稀疏負載存在局部性

通過對稀疏模型進行調研我們發現，稀疏矩陣中非零元素的分佈展現出嚴重的非均勻分佈特性。這種非均勻分佈對稀疏矩陣的計算效率產生了巨大的負面影響：

局部過密：部分區域的非零元素過於密集導致該區域不再適合稀疏矩陣運算；
局部過稀：部分區域的非零元素過於稀疏導致該區域相對於其他區域負載過低，造成計算單元負載失衡問題。這些問題嚴重影響了稀疏算子的執行效率。

圖 2. Llama-2-7B 模型整體 70% 稀疏率場景下非零元素的分佈。第 0、1、2 層的 Query weight 矩陣中不同區域的稀疏率跨度很大（30%~99%），展現出嚴重的非均勻分佈特徵；第 30、31 層的 Query weight 矩陣中不同區域的非均勻分佈特性有所緩解（55%~99%）但依然嚴重。

CROSS：稀疏編譯的破局之道

為應對上述挑戰，CROSS 引入了一套全新的編譯優化流程。CROSS 首先對稀疏矩陣的結構特點進行深入分析，通過代價模型精準判斷稀疏與密集區域的不同計算需求，並自動分配最優的計算資源。其關鍵步驟包括：

1）代價模型構建：首先，我們對不同 block 形狀下、不同稀疏率下的稀疏矩陣乘（SpMM）和密集矩陣乘（GEMM）執行時間進行分析並建立代價模型（block 內的稀疏分佈假設為均勻分佈），如圖 3 所示。SpMM 開銷明顯高於 GEMM 開銷的稀疏率範圍稱為密集區（Dense band）；將 SpMM 開銷明顯低於 GEMM 開銷的稀疏率範圍稱為稀疏區（Sparse band），將 SpMM 與 GEMM 的執行開銷相近的區域稱為搖擺區（Swing band）。

圖 3. 不同稀疏率下 SpMM 與 GEMM 的執行時間分佈（SpMM 由 Sputnik 實現，GEMM 由 cuBlas 實現，矩陣形狀 M=N=K=256, batch=10）。

2） Intra-batch 負載均衡：其次，我們將模型中的稀疏矩陣拆分為多個 block 並依據代價模型評估每個 block 適合的計算範式和計算開銷。然後，我們依據 block 之間是否具有累加關係對整個矩陣的計算開銷建立代價模型，如圖 4（b）所示。針對矩陣中存在的負載不均衡問題，我們將稀疏計算與密集計算分別映射到不同的計算單元執行。當稀疏計算與密集計算的負載差異較大時，我們將搖擺類型的 block 轉換為負載較小的類型，以實現單 batch 稀疏矩陣乘法的計算單元負載均衡（如圖 4（c）所示）。

圖 4. Intra-batch 負載均衡策略。（a）一個稀疏分佈不均勻的稀疏矩陣案例。（b）原始稀疏矩陣乘法的執行開銷。（c）負載均衡策略下的矩陣乘法執行開銷。

3） Inter-batch 負載均衡：此外，由於 batch 之間使用相同的稀疏權值矩陣，當 batch size 較大時，矩陣中不同位置的負載失衡問題持續積累而變得更加嚴重。針對該問題我們將相鄰兩個 batch 之間的負載與計算單元的映射關係進行了重排。如圖 5（a）所示，簡單的將相鄰兩個 batch 合併執行會造成不同位置的負載失衡效應持續積累，造成更嚴重的負載失衡問題。為了應對該問題，我們對不同計算單元的負載進行重排序，相鄰兩個 batch 按照不同的順序進行計算單元映射，以實現 batch 之間的負載均衡。