30行代码，500万长文本推理提速8倍！「树注意力」让GPU越多省的越多

AIGC动态欢迎阅读

原标题：30行代码，500万长文本推理提速8倍！「树注意力」让GPU越多省的越多
关键字：注意力,能量,梯度,函数,拓扑
文章来源：量子位
内容字数：0字

内容摘要：

梦晨发自凹非寺量子位 | 公众号 QbitAI跨GPU的注意力并行，最高提速8倍，支持512万序列长度推理。
环注意力（Ring Attention）后继者——树注意力（Tree Attention）来了。
最关键之处在于，通信步数随设备数量成对数增长，而不是线性增长。
换句话说，树注意力的优势随着设备数量增大会更加明显。实验中，在128卡、512万序列长度设置时达到最高8倍加速。
与环注意力相比，峰值内存占用也能节省不少。
相关代码已经开源，基于谷歌jax框架，已和Flash Attention整合，实现起来只需要30行代码。
论文一公布，就被业界评价为“对高推理需求的大型公司很重要”。
这下和黄仁勋的GPU“买的越多，省的越多”论对上了，英伟达再次赢麻。
注意力机制的能量视角首先简单回顾一下这次被拿来对比的环注意力，由UC伯克利大牛Pieter Abeel团队提出。
环注意力被认为是让上一波大模型纷纷扩展到百万上下文的关键，从谷歌Gemini 1.5到后来的Llama 3.1系列都用了它的某种变体。
简单来说，环注意力的核心思想是将长序列分成多个Block，每个GPU处理一个。

原文链接：30行代码，500万长文本推理提速8倍！「树注意力」让GPU越多省的越多