突破无规则稀疏计算边界，编译框架CROSS数倍提升模型性能

首页 » 突破无规则稀疏计算边界，编译框架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 之间的负载均衡。