提出新型8位低精度softmax方案,显著提升Transformer推理速度。
BAPS: A Fine-Grained Low-Precision Scheme for Softmax in Attention via Block-Aware Precision reScaling
- 采用专用8位浮点格式与分块精度重校准,降低计算开销。
- 数据传输量减半,指数单元面积大幅减少,精度损失小。
- 适合追求高效推理的AI芯片设计与模型部署者。
随着量化矩阵乘法加速收益趋于饱和,Softmax操作成为Transformer推理中的关键瓶颈。该瓶颈源于两大硬件限制:(1) 矩阵与向量计算核心间数据带宽有限;(2) 高精度(FP32/FP16)指数单元(EXP2)面积开销大。为此,本文提出一种新颖的低精度工作流,采用特定8位浮点格式(HiF8)与分块感知精度重校准(block-aware precision rescaling)实现Softmax计算。关键创新在于,无需显著损失模型精度即可实现低精度化。具体而言,(i) 通过将矩阵乘法输出限制在8位,使所需数据移动带宽减半;(ii) 在8位精度下进行指数运算,大幅降低EXP2单元面积。在语言模型与多模态模型上的广泛评估验证了方法有效性。该工作缓解了向量计算瓶颈,可在不增加芯片面积的前提下实现端到端推理吞吐量翻倍,并为未来软硬件协同设计提供可行路径。
原文摘要 · Abstract (English)
As the performance gains from accelerating quantized matrix multiplication plateau, the softmax operation becomes the critical bottleneck in Transformer inference. This bottleneck stems from two hardware limitations: (1) limited data bandwidth between matrix and vector compute cores, and (2) the significant area cost of high-precision (FP32/16) exponentiation units (EXP2). To address these issues, we introduce a novel low-precision workflow that employs a specific 8-bit floating-point format (HiF8) and block-aware precision rescaling for softmax. Crucially, our algorithmic innovations make low-precision softmax feasible without the significant model accuracy loss that hampers direct low-precision approaches. Specifically, our design (i) halves the required data movement bandwidth by enabling matrix multiplication outputs constrained to 8-bit, and (ii) substantially reduces the EXP2 unit area by computing exponentiations in low (8-bit) precision. Extensive evaluation on language models and multi-modal models confirms the validity of our method. By alleviating the vector computation bottleneck, our work paves the way for doubling end-to-end inference throughput without increasing chip area, and offers a concrete co-design path for future low-precision hardware and software.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。