Lightning Attention-2：处理大语言模型中无限序列长度的免费午餐

线性注意力是一种高效的注意力机制，近来已成为传统 softmax 注意力的有前景的替代方案。凭借以线性计算复杂度处理 token 的能力，线性注意力在理论上可以处理无限长度的序列而不牺牲速度，即在固定内存消耗下，对不同序列长度保持恒定的训练速度。然而，由于累积求和（cumsum）的问题，当前的线性注意力算法无法在因果设定中展示其理论优势。本文提出 Lightning Attention-2，这是首个使线性注意力能够实现其理论计算优势的线性注意力实现。为此，我们利用分块思想，分别处理线性注意力计算中的块内和块间分量。具体而言，我们对块内采用传统的注意力计算机制，并对块间应用线性注意力核技巧。在前向和后向过程中均采用分块技术，以充分利用 GPU 硬件。我们在 Triton 中实现算法，使其具有 IO 感知且硬件友好。在不同模型大小和序列长度上进行了各种实验。Lightning Attention-2 无论输入序列长度如何，都保持一致的训练和推理速度，且显著快于其他注意力机制。源代码可在 https://github.com/OpenNLPLab/lightning-attention 获取。