摩爾線程科研團(tuán)隊(duì)近期公布了一項(xiàng)創(chuàng)新研究成果,名為《Round Attention:利用輪次塊稀疏性革新多輪對話優(yōu)化路徑》。這一突破性進(jìn)展使得推理引擎的端到端延遲顯著低于當(dāng)前主流的Flash Attention,同時(shí)在鍵值緩存(kv-cache)顯存占用上實(shí)現(xiàn)了最高82%的節(jié)省。
隨著AI大型語言模型的快速發(fā)展,語言模型服務(wù)在日常問題解決任務(wù)中的應(yīng)用日益廣泛。然而,長時(shí)間的用戶交互帶來了兩大挑戰(zhàn):一是上下文長度的急劇增加導(dǎo)致自注意力機(jī)制的計(jì)算開銷劇增,因其復(fù)雜度與長度的平方成正比;二是鍵值緩存技術(shù)雖然在一定程度上緩解了冗余計(jì)算問題,但隨之而來的GPU內(nèi)存需求激增,限制了推理批處理的規(guī)模,降低了GPU的利用率。
為了應(yīng)對這些挑戰(zhàn),摩爾線程提出了Round Attention機(jī)制。該機(jī)制的核心在于以輪次為單位分析Attention規(guī)律,專為多輪對話場景設(shè)計(jì)。通過對輪次粒度的Attention分布進(jìn)行深入研究,摩爾線程發(fā)現(xiàn)了兩個(gè)重要規(guī)律,這些規(guī)律為優(yōu)化提供了理論基礎(chǔ)。
基于這些發(fā)現(xiàn),摩爾線程進(jìn)一步設(shè)計(jì)了Round Attention推理流水線。這一流水線將稀疏性從傳統(tǒng)的Token級(jí)提升到了塊級(jí),通過選取最相關(guān)的塊參與Attention計(jì)算,顯著減少了計(jì)算耗時(shí)。同時(shí),將不相關(guān)的塊卸載到CPU內(nèi)存,從而有效節(jié)省了顯存占用。這一策略在保持推理精度的前提下,顯著提升了推理效率并降低了資源消耗。
摩爾線程指出,輪次塊稀疏性具有三大顯著優(yōu)勢:首先,以輪次為自然邊界的劃分保證了語義的完整性;其次,在分水嶺層實(shí)現(xiàn)了注意力的穩(wěn)定性;最后,在端到端層面實(shí)現(xiàn)了存儲(chǔ)與傳輸?shù)膬?yōu)化。這些優(yōu)勢共同促成了Round Attention的高效表現(xiàn)。
實(shí)驗(yàn)結(jié)果顯示,與主流的Flash Attention推理引擎相比,Round Attention在端到端延遲方面表現(xiàn)出色,同時(shí)在kv-cache顯存占用上節(jié)省了55%至82%。在主觀評測和客觀評測的兩個(gè)數(shù)據(jù)集上,模型推理準(zhǔn)確率基本保持不變,驗(yàn)證了Round Attention的有效性和實(shí)用性。