量子位
·
30行代码，500万长文本推理提速8倍！「树注意力」让GPU越多省的越多
💡
原文中文，约1900字，阅读约需5分钟。
📝
内容提要
这篇文章介绍了一种名为“树注意力”的新型注意力机制，可以在多个GPU上并行运行，提高长文本推理的速度。树注意力的关键之处在于通信步数随设备数量成对数增长，而不是线性增长，因此随着设备数量增大，其优势更加明显。与之前的环注意力相比，树注意力不仅可以节省峰值内存占用，还能保持完整的注意力计算。树注意力的设计还充分利用了GPU集群的两级拓扑特点，进一步提高了性能。这项研究对于高推理需求的大型公司非常重要。
🎯
关键要点
- 树注意力是一种新型注意力机制，支持在多个GPU上并行运行，最高提速8倍。
- 树注意力的通信步数随设备数量成对数增长，优势随着设备数量增大而更加明显。
- 与环注意力相比，树注意力节省峰值内存占用，并保持完整的注意力计算。
- 树注意力设计利用了GPU集群的两级拓扑特点，提高了性能。
- 相关代码已开源，基于谷歌jax框架，集成Flash Attention，仅需30行代码实现。
- 树注意力团队来自Zyphra，专注于边缘AI和端侧AI，具有数学和理论物理背景。
➡️
