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

在現(xiàn)代 AI 模型的快速迭代中，如何在保持模型精度的同時(shí)提升計(jì)算效率成為關(guān)鍵課題。尤其在大規(guī)模 AI 推理中，非結(jié)構(gòu)化稀疏矩陣的計(jì)算效率低下成為難以突破的瓶頸。面對(duì)這一挑戰(zhàn)，我們自主研發(fā)了 CROSS—— 一種創(chuàng)新的端到端稀疏編譯優(yōu)化方案，為 AI 推理帶來細(xì)粒度稀疏計(jì)算的加速效果。

稀疏計(jì)算的挑戰(zhàn)：如何處理非均勻稀疏分布

非結(jié)構(gòu)化細(xì)粒度稀疏場(chǎng)景下模型推理效率低下問題是 AI 編譯社區(qū)面對(duì)的關(guān)鍵問題之一。相比于密集算子加速庫（cuBlas），主要的稀疏算子加速庫或編譯框架需要在較高稀疏率下才能獲得收益，而過高的稀疏率需求可能使我們面臨模型精度下降的風(fēng)險(xiǎn)。

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

稀疏計(jì)算的機(jī)會(huì)：稀疏負(fù)載存在局部性

通過對(duì)稀疏模型進(jìn)行調(diào)研我們發(fā)現(xiàn)，稀疏矩陣中非零元素的分布展現(xiàn)出嚴(yán)重的非均勻分布特性。這種非均勻分布對(duì)稀疏矩陣的計(jì)算效率產(chǎn)生了巨大的負(fù)面影響：

• 局部過密：部分區(qū)域的非零元素過于密集導(dǎo)致該區(qū)域不再適合稀疏矩陣運(yùn)算；
• 局部過?。翰糠謪^(qū)域的非零元素過于稀疏導(dǎo)致該區(qū)域相對(duì)于其他區(qū)域負(fù)載過低，造成計(jì)算單元負(fù)載失衡問題。這些問題嚴(yán)重影響了稀疏算子的執(zhí)行效率。

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

CROSS：稀疏編譯的破局之道

為應(yīng)對(duì)上述挑戰(zhàn)，CROSS 引入了一套全新的編譯優(yōu)化流程。CROSS 首先對(duì)稀疏矩陣的結(jié)構(gòu)特點(diǎn)進(jìn)行深入分析，通過代價(jià)模型精準(zhǔn)判斷稀疏與密集區(qū)域的不同計(jì)算需求，并自動(dòng)分配最優(yōu)的計(jì)算資源。其關(guān)鍵步驟包括：

1） 代價(jià)模型構(gòu)建：首先，我們對(duì)不同 block 形狀下、不同稀疏率下的稀疏矩陣乘（SpMM）和密集矩陣乘（GEMM）執(zhí)行時(shí)間進(jìn)行分析并建立代價(jià)模型（block 內(nèi)的稀疏分布假設(shè)為均勻分布），如圖 3 所示。SpMM 開銷明顯高于 GEMM 開銷的稀疏率范圍稱為密集區(qū)（Dense band）；將 SpMM 開銷明顯低于 GEMM 開銷的稀疏率范圍稱為稀疏區(qū)（Sparse band），將 SpMM 與 GEMM 的執(zhí)行開銷相近的區(qū)域稱為搖擺區(qū)（Swing band）。

圖 3. 不同稀疏率下 SpMM 與 GEMM 的執(zhí)行時(shí)間分布（SpMM 由 Sputnik 實(shí)現(xiàn)，GEMM 由 cuBlas 實(shí)現(xiàn)，矩陣形狀 M=N=K=256, batch=10）。

2） Intra-batch 負(fù)載均衡：其次，我們將模型中的稀疏矩陣拆分為多個(gè) block 并依據(jù)代價(jià)模型評(píng)估每個(gè) block 適合的計(jì)算范式和計(jì)算開銷。然后，我們依據(jù) block 之間是否具有累加關(guān)系對(duì)整個(gè)矩陣的計(jì)算開銷建立代價(jià)模型，如圖 4（b）所示。針對(duì)矩陣中存在的負(fù)載不均衡問題，我們將稀疏計(jì)算與密集計(jì)算分別映射到不同的計(jì)算單元執(zhí)行。當(dāng)稀疏計(jì)算與密集計(jì)算的負(fù)載差異較大時(shí)，我們將搖擺類型的 block 轉(zhuǎn)換為負(fù)載較小的類型，以實(shí)現(xiàn)單 batch 稀疏矩陣乘法的計(jì)算單元負(fù)載均衡（如圖 4（c）所示）。

圖 4. Intra-batch 負(fù)載均衡策略。（a）一個(gè)稀疏分布不均勻的稀疏矩陣案例。（b）原始稀疏矩陣乘法的執(zhí)行開銷。（c）負(fù)載均衡策略下的矩陣乘法執(zhí)行開銷。

3） Inter-batch 負(fù)載均衡：此外，由于 batch 之間使用相同的稀疏權(quán)值矩陣，當(dāng) batch size 較大時(shí)，矩陣中不同位置的負(fù)載失衡問題持續(xù)積累而變得更加嚴(yán)重。針對(duì)該問題我們將相鄰兩個(gè) batch 之間的負(fù)載與計(jì)算單元的映射關(guān)系進(jìn)行了重排。如圖 5（a）所示，簡(jiǎn)單的將相鄰兩個(gè) batch 合并執(zhí)行會(huì)造成不同位置的負(fù)載失衡效應(yīng)持續(xù)積累，造成更嚴(yán)重的負(fù)載失衡問題。為了應(yīng)對(duì)該問題，我們對(duì)不同計(jì)算單元的負(fù)載進(jìn)行重排序，相鄰兩個(gè) batch 按照不同的順序進(jìn)行計(jì)算單元映射，以實(shí)現(xiàn) batch 之間的負(fù)載均衡。

圖 5. Inter-batch 負(fù)載均衡策略。（a）簡(jiǎn)單粗暴的將相鄰兩個(gè) batch 合并會(huì)造成負(fù)載失衡效應(yīng)累積。（b）對(duì)相鄰 batch 中不同計(jì)算單元的負(fù)載重排能大幅緩解負(fù)載均衡問題。

實(shí)驗(yàn)成果：顯著的性能提升

實(shí)驗(yàn)結(jié)果表明，相比于其他稀疏矩陣加速庫或編譯框架，CROSS 在不同稀疏率下都獲得了顯著性能提升，與業(yè)界最優(yōu)設(shè)計(jì)相比平均獲得 2.03× 的性能提升。相比于密集計(jì)算（cuBlas）,CROSS 在稀疏率超過 60% 時(shí)開始獲得正收益，顯著突破了傳統(tǒng)無規(guī)則稀疏加速設(shè)計(jì)的收益邊界。

圖 6. 五種整體模型稀疏率下，不同稀疏加速設(shè)計(jì)相比于密集加速庫（cuBlas）的模型推理性能。相比于 cuBlas，我們最高可以獲得 3.75× 性能收益，同時(shí)我們?cè)谙∈杪食^ 60% 時(shí)開始獲得正收益，而其他方案則需要接近或超過 80% 稀疏率。

CROSS 的未來：推動(dòng)稀疏編譯應(yīng)用普及

CROSS 的成功不僅在于提升了稀疏矩陣計(jì)算的效率，更為未來 AI 推理在稀疏計(jì)算場(chǎng)景下的廣泛應(yīng)用奠定了堅(jiān)實(shí)的基礎(chǔ)。在 AI 模型規(guī)模不斷擴(kuò)展的今天，稀疏性在大模型中廣泛存在，CROSS 為稀疏編譯提供了高效、靈活、可持續(xù)的發(fā)展路徑，助力未來 AI 應(yīng)用的高效部署。