研究摘要
基于Transformer的模型在多个领域取得了显著的进步，这主要得益于自注意力机制能够捕捉输入序列中的上下文关系。然而，处理长序列对于Transformer模型来说仍然计算成本高昂，这主要是由于自注意力机制带来的O(n²)复杂度。为了解决这一问题，稀疏注意力被提出，旨在将二次依赖降低为线性。尽管如此，高效部署稀疏Transformer仍面临两大主要障碍：1）由于算法的近似特性导致注意力分散，现有的系统优化对稀疏Transformer效果不佳；2）输入序列的变异性导致计算和内存访问效率低下。本文提出创新性Transformer框架Raptor-T，致力于解决现有稀疏注意力模型在长序列与变长序列处理中的系统性挑战。针对传统Transformer模型因自注意力机制O(n²)复杂度导致的长序列处理效率瓶颈，本工作通过算法与系统协同优化实现三重突破：首先，设计融合式内存高效多头注意力机制，有效解决稀疏注意力计算碎片化问题；其次，创新性引入异步数据处理架构，消除GPU阻塞操作，实现计算资源的高效利用；最后，提出动态填充优化策略，通过智能序列打包技术降低填充冗余度，显著提升GPU计算均衡性。
研究背景
图1： 稀疏注意力计算模型
Transformer模型架构由编码器和解码器组成。编码器包含一个自注意力层和一个前馈神经网络，能够捕捉当前信息和上下文信息，从而更深入地理解输入序列。解码器具有类似的层结构，并额外包含一个注意力层，用于强调关键信息，从而增强模型的预测能力。注意力机制使得模型能够自动学习并更好地处理长距离依赖关系。
然而，Transformer中的全注意力机制存在对序列长度的二次依赖问题，在处理长序列时会占用大量内存和计算资源。为了解决这一问题，研究人员提出了稀疏Transformer模型，例如BigBird。这类模型通过计算有限的相似度得分来生成稀疏矩阵，从而减少计算和内存需求。BigBird在自然语言处理任务（如问答和文本摘要）中实现了最先进的性能，同时具备图灵完备性。它将输入张量根据块大小划分为多个块，并计算这些块之间的注意力。多种注意力机制被结合起来，形成了最终的注意力机制，其结构由计算所需的注意力矩阵如图1表示。
尽管BigBird减少了处理长序列所需的计算和内存资源，但其端到端的性能仍面临一些限制。
研究现状
针对于Transformer模型的优化，国内外学者提出了一系列方法。如ByteTransformer和FlashAttention，提出了融合注意力内核以加速计算。然而，这些方法是在稀疏注意力场景下的仍然面临瓶颈。
1. GPU硬件资源浪费
我们发现BigBird花费了近50%的时间执行CPU操作，导致宝贵的GPU资源被阻塞。GPU主要在执行注意力计算之前被阻塞，因为它需要等待CPU为分散的注意力计算准备数据。
2. 更多的CUDA内核
尽管稀疏注意力将计算从二次依赖降低到线性，但BigBird的执行时间仍然比预期要长。主要是由于现有方法需要将稀疏注意力计算分散成多个全注意力计算，从而导致BigBird比Bert包含了更多的计算CUDA内核，进而使得大量时间花费在从全局内存加载和存储数据上。
3. 实际推理场景下的变长的长序列输入
实际场景下的模型输入长度是随机的，从而导致模型会接收到不同长度的序列输入。而现有方法如TurboTransformer和ByteTransformer提供了改进，但是在长序列输入场景下仍然不能避免填充开销，导致性能无法显著提升。
研究方法
图2： Raptor-T 中融合注意力的线程块视图
为了解决这些挑战，我们推出了 Raptor-T，这是一个专为长序列和可变长度输入序列设计的尖端Transformer框架。我们的方法基于一个关键洞察：识别出在处理长序列时，由于稀疏注意力机制导致的分散计算性能瓶颈，以及可变长度输入引起的冗余计算问题。我们没有分别解决这两个问题，而是采取了一种整体方法，协同解决它们。
Raptor-T 包含两个关键机制，旨在有效优化稀疏Transformer。首先，我们引入了一种新颖的融合且内存高效的稀疏多头注意力（MHA）机制，专门为高效计算任意长度序列而设计。通过这种方法，我们减少了不必要的开销并加速了计算，这对于处理长序列尤其有益。其次，我们实现了一种异步机制，以最小化GPU阻塞时间，从而显著提高整体系统效率。这一优化确保了计算资源的更好利用，并提升了框架的性能。以上所有方法均实现了数学上等价的内核，确保了统计结果的一致性。
基于这些机制，我们开发了两种创新算法来处理可变长度序列输入。最小填充算法采用细粒度的块填充，避免了不必要的最大长度填充，从而进一步优化效率，尤其是在处理长序列时。此外，注意力平衡算法根据线程块的执行时间动态分配GPU内的任务，有效解决了计算不平衡问题，并进一步提升了GPU效率。通过整合这些机制和算法，Raptor-T 实现了一个真正全面且高效的解决方案，在处理长序列和可变长度序列时表现出色，同时提供了卓越的性能和资源优化能力。
研究成果
图3： 端到端推理时间对比
我们在NVIDIA A100 GPU上使用长序列和可变长度序列输入对Raptor-T进行了广泛评估。实验结果表明，Raptor-T在端到端性能上优于现有的最先进实现。具体而言，Raptor-T在性能上平均比FlashAttention-2 和 FasterTransformer分别提升了3.41倍和3.71倍。
相关研究成果已收录于 CCF-A 类顶级期刊IEEE TC（Transactions on Computers）上，论文标题《Raptor-T: A Fused and Memory-Efficient Sparse Transformer for Long and Variable-Length Sequences》。