行代码万长文本推理提速倍！树注意力让越多省得越多

行代码，万长文本推理提速倍！树注意力让越多省得越多

梦晨发自凹非寺量子位|公众号QbitAI跨GPU的注意力并行，最高提速8倍，支持512万序列长度推理。环注意力（RingAttention）后继者——树注意力（TreeAttention）来了。最关键之处在于，通信步数随设备数量成对数增长，而不是线性增长。换句话说，树注意力的优势随着设备数量增大会更加明显。实验中，在128卡、512万序列长度设置时达到最高8倍加速。与环注意力相比，峰值内存占用也能节省不少相关代码已经开源，基于谷歌jax框架，已和FlashAttention整合，实现起来只需要30行代码论文一公布，...