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

文章预览

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