arXiv:2503.01873cs.LGcs.AI2025-03

提出PASA算法,让低精度大模型推理更稳更快。

Online Pseudo-average Shifting Attention(PASA) for Robust Low-precision LLM Inference: Algorithms and Numerical Analysis

  • 在线伪均值偏移+全局恢复,避免半精度计算溢出
  • 在昇腾NPU上加速推理,减少数据搬移并提升算力利用率
  • 适用于长序列生成任务,尤其对视频与语言模型有效

长序列生成任务中注意力计算极为耗时。为此,我们基于Flash Attention设计了一种低精度、数学等价的算法PASA,引入在线伪均值偏移和全局恢复两项新技术,实现全程半精度计算而无溢出不稳定或精度显著下降问题。该算法在内存受限的AI硬件(如昇腾NPU)上通过减少数据移动、增加计算浮点操作数,显著提升性能。实验表明,大模型中注意力输入数据的偏差与幅度是导致半精度溢出(>65504)的关键因素,分别源于序列维度的大偏置及Stable-Video-Diffusion模型中查询与键在头维度的共振机制——即相位重合或180度反向,会显著放大注意力分数矩阵元素。此现象同样存在于Qwen2-7B模型中。数值精度通过均方根误差(RMSE)评估,并与高精度注意力生成结果对比验证。

原文摘要 · Abstract (English)

Attention calculation is extremely time-consuming for long-sequence inference tasks, such as text or image/video generation, in large models. To accelerate this process, we developed a low-precision, mathematically-equivalent algorithm called PASA, based on Flash Attention. PASA introduces two novel techniques: online pseudo-average shifting and global recovering. These techniques enable the use of half-precision computation throughout the Flash Attention process without incurring overflow instability or unacceptable numerical accuracy loss. This algorithm enhances performance on memory-restricted AI hardware architectures, such as the Ascend Neural-network Processing Unit(NPU), by reducing data movement and increasing computational FLOPs. The algorithm is validated using both designed random benchmarks and real large models. We find that the large bias and amplitude of attention input data are critical factors contributing to numerical overflow ($>65504$ for half precision) in two different categories of large models (Qwen2-7B language models and Stable-Video-Diffusion multi-modal models). Specifically, overflow arises due to the large bias in the sequence dimension and the resonance mechanism between the query and key in the head dimension of the Stable-Video-Diffusion models. The resonance mechanism is defined as phase coincidence or 180-degree phase shift between query and key matrices. It will remarkably amplify the element values of attention score matrix. This issue also applies to the Qwen models. Additionally, numerical accuracy is assessed through root mean square error (RMSE) and by comparing the final generated texts and videos to those produced using high-precision attention.

大模型推理低精度计算注意力机制昇腾NPU

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。