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越多省的越多