Tree Attention: Topology-Aware Decoding for Long-Context Attention on GPU Clusters

本文是LLM系列文章，针对《Tree Attention: Topology-Aware Decoding for Long-Context Attention on GPU Clusters》的翻译。

摘要

自注意力是现代Transformer架构的核心数学运算，由于其序列长度的二次复杂性，也是一个重要的计算瓶颈。在这项工作中，我们推导了标量能量函数，其梯度计算了自我注意块，从而阐明了自注意力的理论基础，提供了对操作的贝叶斯解释，并将其与Hopfield网络等基于能量的模型紧密联系起来。我们的公式表明，通过树约简可以并行有效地计算序列轴上的约简。我们的算法，用于在多个GPU上并行化注意力计算，使跨设备解码的执行速度比Ring attention等替代方法快得多（在我们的实验中快8倍），同时还需要更少的通信量，减少2倍的峰值内存。我们的代码可在此处公开获取。

1 引言

2 相关工作

3 自注意力

4 自注意力作为梯度

5 树注意力

6 结果

7 讨论和结论

本文推导了自注意力的能量函数，并证明了该函数导数的计算如何为并行计算注意提供了一种新颖有效的方法。当在多个设备上执行解码时，这一优势尤其明显，在这种情况下，我们的树注意力使我们能够通过渐进优越的算法大大优于环注意力，当我们在5.12M的序列长度上使用128个GPU时，速度提高了8倍。我们还看到，我们使用的AllReduce操作涉及发送部分减少的