BAPS: 一种针对注意力机制中softmax的细粒度低精度方案 via 块意识精度重标定
BAPS: A Fine-Grained Low-Precision Scheme for Softmax in Attention via Block-Aware Precision reScaling
AI总结:
BAPS提出了一种细粒度低精度方案,通过块意识精度重标定和HiF8格式实现高效softmax,减少数据带宽和指数运算单元面积,提升Transformer推理吞吐量。
摘要:
随着加速量化矩阵乘法的性能增益趋于平缓,softmax操作成为Transformer推理中的关键瓶颈。这一瓶颈源于两个硬件限制:(1)矩阵和向量计算核心之间的数据带宽有限;(2)高精度(FP32/16)指数运算单元(EXP2)的显著面积成本。为了解决这些问题,我们引入了一种新的低精度工作流程,采用特定的8位浮点格式(HiF8)和块意识精度重标定进行softmax。关键在于我们的算法创新使低精度softmax成为可能,而不会像直接低精度方法那样显著损失模型精度。具体而言,我们的设计(i)通过将矩阵乘法输出限制为8位,将所需的数据带宽减少了一半;(ii)通过在低(8位)精度下计算指数运算,大幅减少了EXP2单元的面积成本。对语言模型和多模态模型的广泛评估证实了我们方法的有效性。通过缓解向量计算瓶颈,我们的工作为在不增加芯片面积的情况下将端到端推理吞吐量翻倍铺平了道路,并为未来的低精度硬件和软件提供了具体的协同设计路径。
英文摘要:
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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。