该文档属于类型a,即报告了一项原创性研究,以下是针对该研究的学术报告:
Sanger:一种基于可重构架构实现稀疏注意力机制的软硬件协同设计框架
作者及单位
本研究的核心团队来自北京大学信息科学技术学院和中国科学院软件研究所,包括Lu Liqiang*、Jin Yicheng*、Bi Hangrui、Luo Zizhang、Li Peng、Wang Tao以及通讯作者Liang Yun。研究成果发表于2021年10月的MICRO ‘21: 54th Annual IEEE/ACM International Symposium on Microarchitecture,会议地点为希腊线上会议。
学术背景
研究领域与动机
自然语言处理(NLP)和计算机视觉领域中,基于注意力机制(attention mechanism)的模型(如Transformer)通过捕捉全序列上下文信息取得了突破性表现。然而,传统注意力机制存在大量冗余计算,尤其是其核心操作——注意力矩阵(attention matrix)的稠密矩阵乘法(GEMM)会随序列长度呈平方级增长计算复杂度。例如,处理16k tokens的输入时,单个BERT模块的计算量高达861.9 GFLOPs。为缓解计算负担,稀疏注意力(sparse attention)成为研究热点,其通过采样稠密-稠密矩阵乘法(SDDMM)和稀疏-稠密矩阵乘法(SPMM)减少无效运算。但现有方法面临两大挑战:
1. 静态稀疏(static sparsity)虽硬件友好但牺牲计算效率(如块稀疏仅压缩30%-50%计算量);
2. 动态稀疏(dynamic sparsity)虽能适配输入数据特性,但非结构化稀疏模式导致硬件负载不均衡和访存局部性差。
研究目标
提出Sanger框架,通过软硬件协同设计实现动态稀疏注意力的高效加速,核心目标包括:
- 设计动态细粒度结构化剪枝算法,平衡稀疏性与硬件效率;
- 开发可重构硬件架构,支持灵活稀疏模式并消除解码开销;
- 在BERT/GPT-2等模型上实现高压缩比(≥80%稀疏度)且无精度损失。
研究流程与方法
1. 软件层:动态结构化剪枝算法
步骤① 低比特注意力矩阵预测
- 输入:稠密查询矩阵(Q)和键矩阵(K),通过4位对称线性量化(symmetric linear quantization)降低计算开销。
- 操作:量化后矩阵经低精度GEMM和Softmax生成近似注意力矩阵(\hat{S}),计算量仅为全精度版本的1/16。
- 创新点:训练中采用直通估计器(STE)模拟量化误差,保证模型适应性。
步骤② 二值化阈值生成稀疏掩码
- 对(\hat{S})应用全局阈值(t)(如SQuAD任务设定(t=2e-3)),生成二元掩码(M)(元素≥(t)为1,否则为0)。实验表明,该简单方法在BERT中可实现0.08-0.27的稀疏度。
步骤③ 掩码结构化编码
- 问题:非结构化掩码导致硬件并行度低。
- 解决方案:将掩码按PE阵列宽度分块,通过Packing跳过全零行、Splitting拆分过密行,最终生成负载均衡的块结构(如每行非零数≤25%)。例如,CLOTH任务(序列长512)的掩码编码后PE利用率从0.45提升至0.63。
2. 硬件层:可重构脉动阵列设计
核心创新:分数驻留数据流(Score-Stationary Dataflow)
- 统一SDDMM与SPMM:将稀疏分数矩阵(S)固定在PE内,避免反复解码稀疏模式。
- 动态配置PE连接:通过多路选择器(MUX)按掩码动态分配Q/K/V数据路径。
- 可重构PE设计:支持迭代累加(SDDMM阶段)与前向累加(SPMM阶段),通过气泡控制器(Bubble Controller)调节数据传输延迟。
硬件实现
- 工艺:UMC 55nm,频率500MHz,面积16.9mm²,功耗2.76W。
- 关键模块:4位乘法阵列(预处理器)、16×64可重构PE阵列(支持128GB/s HBM带宽)、指数运算单元(基于LUT近似)。
主要结果
1. 模型压缩效果
在GLUE、SQuAD和CLOTH任务上对比静态稀疏方法(如BigBird、LongFormer):
- 精度:BERT稀疏至0.127(SQuAD)时,EM/F1仅下降0.4⁄0.2,优于BigBird(下降1.3⁄1.0);
- 稀疏度:GPT-2和BART分别达0.15-0.35和0.23-0.54,精度损失<0.5%。
2. 硬件加速性能
- 对比GPU/CPU:在BERT上较V100 GPU(FP16)提速4.64倍,能效提升35倍;较AMD Ryzen CPU提速22.7倍。
- 对比专用加速器:超越A3(2.39倍)和Spatten(1.47倍),主要归因于动态稀疏支持与统一数据流设计。
3. 关键数据
- 硬件利用率:编码后PE利用率平均提升1.5倍(如MNLI任务从0.52至0.72);
- 吞吐量:有效吞吐达529 GOP/s(A3为221 GOP/s),主要得益于分数驻留数据流减少60%解码开销。
结论与价值
科学意义
- 提出首个支持动态稀疏注意力的软硬件协同设计框架,解决了非结构化稀疏与硬件效率间的矛盾。
- 理论贡献包括:低比特预测剪枝理论、分数驻留数据流形式化模型、可重构PE的调度优化方法。
应用价值
- 为长序列处理(如文档翻译、视频理解)提供可部署解决方案,实测在16k tokens场景下能耗降低48倍。
- 开源硬件设计(Chisel实现)可扩展至其他稀疏计算场景(如GNN、推荐系统)。
研究亮点
1. 动态稀疏模式创新:通过量化预测+阈值化实现输入自适应剪枝,稀疏度较静态方法提升2-3倍;
2. 硬件架构突破:可重构脉动阵列支持任意非零分布,面积效率较固定架构(如Spatten)提高1.8倍;
3. 端到端贡献:从算法(4位预测)、编译(掩码编码)到硬件(PE动态配置)的全栈优化。
其他有价值内容
- 跨模型泛化性:框架在GPT-2(生成任务)和BART(Seq2Seq)上验证有效性,表明其普适性;
- 能耗分析:HBM带宽优化使能量效率达113倍(vs CPU),满足边缘计算需求。
(注:原文实验部分包含11项任务、3类模型对比及5种基线方法,此处因篇幅限制仅摘录核心结论。完整数据可参考论文Table 3及Figure 8。)