30行代码，500万长文本推理提速8倍！「树注意力」让GPU越多省的越多
内容提要
这篇文章介绍了一种名为“树注意力”的新型注意力机制，可以在多个GPU上并行运行，提高长文本推理的速度。树注意力的关键之处在于通信步数随设备数量成对数增长，而不是线性增长，因此随着设备数量增大，其优势更加明显。与之前的环注意力相比，树注意力不仅可以节省峰值内存占用，还能保持完整的注意力计算。树注意力的设计还充分利用了GPU集群的两级拓扑特点，进一步提高了性能。这项研究对于高推理需求的大型公司非常重要。
延伸解读
树注意力与环注意力的关键差异
树注意力与环注意力都采用分块计算和跨设备并行，但通信模式不同。环注意力中，GPU排成圆环传递Key-Value信息，通信步数随设备数量线性增长；而树注意力通过树形Allreduce汇总局部能量函数，通信步数随设备数量对数增长。因此，设备越多，树注意力的通信优势越明显，在128卡、512万序列长度下最高提速8倍。此外，树注意力还能节省峰值内存，并保持完整的注意力计算。
能量视角如何支撑树注意力设计
树注意力团队将自注意力解释为能量函数对辅助变量的梯度，即注意力结果等于标量能量函数F对ζ的梯度在ζ=0处的值。由于能量函数中的logsumexp和max运算满足结合律，可以按任意顺序计算。基于此，团队先在各GPU上并行计算局部能量函数，再通过树状Allreduce汇总，最后用自动微分取梯度得到注意力输出。这一过程仅需与计算能量函数相同的时间开销，显存占用也几乎没有额外负担。
GPU集群拓扑与树注意力的适配性
树注意力在设计上充分利用了GPU集群的两级拓扑特点：同节点内使用高速NVLink，节点间则依赖IB或以太网。这种拓扑下，树注意力的通信模式能更好地将通信与计算重叠，避免被最慢的互联带宽制约。相比之下，环形注意力天然不适应这种拓扑，难以实现良好的通信与计算重叠，最终性能受限于最慢的互联带宽。因此，树注意力在具有两级拓扑的GPU集群上能发挥更大优势。
实现与未来优化空间
树注意力的相关代码已开源，基于谷歌jax框架，并已与Flash Attention整合，实现仅需30行代码。这降低了研究者和开发者尝试该技术的门槛。目前，单GPU内部虽理论上可用类似策略提速，但受限于流式处理器（SM）间共享内存的通信方式，优势并不明显。不过，英伟达在H100上实验性地支持了SM间点对点指令，这为未来单卡注意力优化带来了新的想象空间。
Q&A
树注意力是什么？
树注意力是一种新型注意力机制，支持在多个GPU上并行运行，能够提高长文本推理的速度，最高可提速8倍。
树注意力相比于环注意力有什么优势？
树注意力的通信步数随设备数量成对数增长，节省峰值内存占用，并保持完整的注意力计算，优势在设备数量增大时更加明显。
如何实现树注意力？
树注意力的相关代码已开源，基于谷歌jax框架，集成Flash Attention，仅需30行代码即可实现。
树注意力的设计如何利用GPU集群的特点？
树注意力设计充分利用了GPU集群的两级拓扑特点，在同节点内使用高速NVLink，节点间依赖IB或以太网等，提高了性能。
树注意力对大型公司的意义是什么？
树注意力的研究被认为对高推理需求的大型公司非常重要，能够显著提升长文本推理的效率。
树注意力团队的背景是什么？
树注意力团队来自Zyphra，专注于边缘AI和端侧AI，团队成员具有数学和理论物理背景。
