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

AIxiv专栏是本站发布学术、技术内容的栏目。过去数年，本站AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本篇工作已被 HPCA 2025 接收，由上海交大先进计算机体系结构实验室蒋力3 g y D D教授课题组（IMPACT）完成，同时也获得H h f $ + % U了上海期智研究院的支持。第一作者是刘方鑫老师与博士生黄世远。

在现代 AI 模型的快速迭代中，如何在保持模型精度的同时提升计算效率成V h g v i为关键课题。尤其在大规模 AI 推理中，非结构化稀疏矩阵的计算效率低下成为难以突破的瓶p / ] q . k m p颈。面对这一挑战，我们自主研发了 C` t \ 0 @ vROSS—— 一种创新的端到端稀疏编译优化方案，为 AI 推理带来细粒度稀疏计算的加速效果。

稀疏计算的挑战：如何处理非均匀稀疏分布

非结构化细粒度稀疏场景下模型推理效率低下问题是 AI 编译社区面对的关键问题之一。相比于密集算子加速库（cuBlas），P I U G主/ h X $ &要的稀疏算子加速库或编译W = D框架需要在较高稀疏率下才能获得收益，而过高的稀疏率需求可能使我们面临模型精度下降的风险。

^{图 1. 相1 g * 0 7 –比于 cuBlas，不同稀疏加速库或编译框架在不同稀疏率下的加速比。Sputnik、TVM-Sparse、SparseTIR、ASpT 和 cuSPARSE 在稀J ; } z y ; d 4疏率超过 76%、80.5%、82.6%、89.4% 和 98.1% 时才能获得正向收益（稀疏矩阵源于 Bert 模型中的稀疏权值矩阵）。}

稀疏计算的机会：稀疏负载存在局部性

通过对稀疏模型进行调研D j C我们发现，稀疏矩阵中非零元素的分/ 7 U M l ( \布展现出严重的非均匀分布特性。这5 F \ y Y种非均匀分布对稀疏矩阵的计算效率产生了巨大的负面影响：

局部过密：部分区域的非零元素过于密集导致该区域不再适合稀疏矩阵运算；
局部过稀：部分区域的非零元素过于稀疏导致该区域相对于其他区域负载过低，Z R ` C Q e T造成计算单元负载失衡问题。这些问题严重影响了稀疏算子的执行效率。

^{图 2. Llama-2-7Bn o W p @ # ; 模型整体 70% 稀疏率场景下非零元素的分布。第 0、1、2 层的 Query weight 矩阵中不同区域的稀疏率跨度很大（30%~99%），展现出严重的非均匀分布特征；第 30、3` ] , \ 41 层的 Query wea 5 J ` 5 L c p iight 矩阵中不同区域的非均匀分布特性有所缓解（55%~99M T W d o D z N%）但依然严重。}

CROSS：稀疏编译的破局之道

为应对上述挑战，CROSS 引入了一套全新的编译优化k J L Z _ 4流程。CROSS 首先对稀0 N ( w ?疏矩阵的结构特点进行深入分析，通过代价模型精准判断稀疏与密集区域的不同计& P L $ \ } t { 6算需求，并自动分配最优的计算资源。其关键步骤包括：

1）代价模型构建：首先，我们对不同 block 形状下、不同稀疏率下的稀疏矩阵乘（SpMM）和密集矩阵乘（GEMM）执行时间v a ` Q m [ i q进行分析并建立代价模型（block 内的稀疏分布假设为均匀分布），如图 3 所示。SpMM 开销明显高于 GEMM 开销的稀疏率范围称为密集区（Dense band）；将 SpMM 开销明显低于 GEMMo = L c N 4 = 开销的稀疏率范围称为稀疏区（Spary 9 y q ?se band），将 SpMM 与 GEMM 的执行开销相近的区域& W C \ p { U ~ E称为摇摆区（Swing band）。

^{图 3. 不同+ T c N稀疏率下 SpMM 与 GEMM 的执行时间分布（SpMM 由 Sputnik 实现，GEMM 由 cuBlas 实} O # d ` & V I H现，矩阵形状 M=N=K=256, batch=10）。}

2G C d U K） Intra-batch 负载均衡：其次，我们将模型中的稀疏矩阵拆分为多个 block 并依据代价模型评估每个 block 适合[ k q K W y 1 s的计算范式和计算开销。然后，我们依据 block 之间是否具有累加关系对整个矩阵的计算开销建立代价模型，如图 4（b）所示。针o X s ~ 7对矩阵中存在的负载不均衡问题，我们将稀疏计算与密集计算? G ; , ` e *分别映射到不同的计算单元执行。当稀疏计算与密集计算的负载差异较大时，我们将摇摆类型的 block 转换为负载较小的类型，以实现单 batch 稀* x 9 !疏矩阵乘法的计算单元负载均. 4 W衡（如图 4（c）所示）。

^{图 4. Intra-batch 负载均衡策略。（a）一个稀疏分布不均匀的稀疏矩阵案例。（b）原始稀疏矩阵乘法的执行开销。（Q ] * ? 6 b # C ic）负载均衡策略下的矩阵乘法执行开销。}

3） Inter-bato @ d % 1 h 1 o xch 负载均衡：此外，由于 batch 之间使用相同的稀疏权值矩阵，当 batch size 较大时，矩阵中不同位置的负载失衡问题持续积累而变得更加严重。针对该问题我们将相邻两个 batch 之间的负载与计算单元的映射关系进行了重排。如图 5（a）所3 g o U R E示，简单的将相邻两个 batch 合并执行会+ 1 b造成不同位置的负载失衡效应持续积累，造成更% l N严重的负载失衡问题。为了应对该问题，我们对不同计算单元的负载进行重排序，相邻两个 batch 按照不同的顺序进行计算单元映射，以实现 batchV 4 7 E k Q 3 . L 之间的负载均衡。

^{图 5. Inter-batct . = F ; ^ b I 1h 负载均衡策略。（a）简单粗暴的将相邻两个 batch 合并会造成负载失衡效应累积。（b）对相邻 batch 中不同计算单元的负载重排能大幅缓解负载均衡问题。}

实验成果：显著的性能提升

实验结果表明，相Z k ( O x u v 4 b比于其他稀疏矩阵加速库或编译框架，CROSS 在不同稀疏率下都获得了显著性能提升，与业界最优设计5 v Y 5相比平均获得 2.03 的性能提升。相比于密集计算（cuBlas）,CROSS 在稀疏率超过 60% 时开始获得正收r : 4 : 3益，显著突破了传统无规则稀疏加s R ) } o V速设计的收益边界。

^{图 6. 五种整体模型稀疏率下，不同稀疏加速设计相比于密集加F ^ T K ) / p Y .速库（cuBlas）的模型推理性能。相比于 cuBlas，我们最高可以获得 3.75 性能收益，同时我们在稀疏率超过 60% 时开始获得正收益，而其他方案则需要接近或超B v ` c * V k过 80% 稀疏率。}

CROSS 的未来：推动稀疏编H B G q R译应用普及

CROSS 的成功不仅在于提X @ 3 x k u 5升了稀疏矩阵计算的效率，更为未来 A~ \ uI 推理在稀疏计算场景下的广泛应用奠定了坚实的基础。在 AI 模型规模不断扩展的今天，稀疏性S * :在大模型中广泛存在，CROSS 为稀疏编译提供了高效、灵活、可持续的f + . * : l发展路H $ p P Z z径，助力未来 AI 应用的高效部署。

以上就是突破无规则稀疏计算边界，编译框架CROSS数倍提升模型性能的详细内容！

百度未收录

本文由【好易之】整理自网络！
原创文章，作者:【好易之】如转载请注明出处：https://www.zhengjiaxi.com/zxwd/itzx/118750.html
如有侵权，请邮件联系 aoduoye@qq.com 删除。
本站发布的文章及附件仅限用于学习和研究目的;不得将上述内容用于商业或非法用途，否则后果请用户自负。
本站信息来自网络，版权争议与本站无关。您必须在下载后的24个小时之内，从您的电脑中彻底删除上述内容。
如果您喜欢该程序，请支持正版软件，购买注册，得到更好的正版服务。如有侵权请邮件与我们联系处理。