Lean Attention: Hardware-Aware Scalable Attention Mechanism for the Decode-Phase of Transformers

基于Transformer的模型已成为自然语言处理、自然语言生成和图像生成等领域中最常用的架构之一。最先进的模型大小不断增加，达到了数十亿个参数。这些巨大的模型需要大量内存，即使在最先进的AI加速器（如GPU）上，也需要承受显著的推断延迟。具体而言，注意力操作的时间和内存复杂度与总上下文长度（即提示和输出令牌）的平方成正比。因此，为了满足依赖于这些大型模型的应用的低延迟需求，提出了多种优化技术，如键值张量缓存和FlashAttention计算。然而，这些技术并不能满足推断过程中不同阶段的计算差异性。因此，我们提出了LeanAttention，一种可扩展的技术，用于计算仅解码器Transformer模型的令牌生成阶段（解码阶段）的自注意力。通过重新设计解码阶段的执行流程，LeanAttention实现了自注意力机制的可扩展性，使其能够应对长上下文长度的挑战。我们发现，在线softmax的关联性质可以被视为一种约简操作，从而使我们能够在这些大上下文长度上并行计算注意力。我们将瓦片式计算的“流K”样式约简扩展到自注意力，从而实现了并行计算，使注意力执行速度平均提高了2.6倍，FlashAttention-2的速度提高了8.33倍，适用于512k上下文长度。

论文旨在解决transformer模型中attention机制计算复杂度高、内存占用大、推理时间长的问题，提出了一种针对decoder-only transformer模型decode-phase阶段的可扩展的self-attention计算技术。

LeanAttention通过重新设计decode-phase的执行流程，将在线softmax的关联性质视为一个约减操作，从而使我们能够将attention计算并行化，从而实现对长上下文长度的attention机制的扩展。

论文使用了多个数据集进行实验，并且与现有的优化技术进行了比较，结果表明LeanAttention的平均attention执行速度比FlashAttention-2提高了2.6倍，对于512k上下文长度可以提高高达8.33倍的速度。值得注意的是，这种技术只适用于decoder-only transformer模型的decode-phase阶段。

在最近的相关研究中，也有一些关于transformer模型的优化技术，比如key-value tensor caching和FlashAttention计算等。