APB – 清華聯(lián)合騰訊等機(jī)構(gòu)推出的分布式長(zhǎng)上下文推理框架

APB是什么

APB（Accelerating Distributed Long-Context Inference by Passing Compressed Context Blocks across GPUs）是由清華大學(xué)等機(jī)構(gòu)共同研發(fā)的一種分布式長(zhǎng)上下文推理框架。它通過(guò)引入稀疏注意力機(jī)制和序列并行推理的方法，有效克服了大型模型在處理長(zhǎng)文本時(shí)所面臨的效率挑戰(zhàn)。APB利用更小的Anchor block和Passing block，并結(jié)合查詢感知的上下文壓縮技術(shù)，減少了計(jì)算開(kāi)銷，同時(shí)精確傳遞關(guān)鍵信息，從而實(shí)現(xiàn)對(duì)長(zhǎng)距離語(yǔ)義依賴的高效處理。在處理128K文本時(shí)，APB的推理速度比Flash Attention快約10倍，比英偉達(dá)的Star Attention快1.6倍，且表現(xiàn)優(yōu)異。APB具備卓越的兼容性，能夠適應(yīng)不同的分布式設(shè)置和模型規(guī)模。

APB的主要功能

• 加速長(zhǎng)上下文推理：APB通過(guò)多主機(jī)近似注意力機(jī)制顯著提升推理速度，相比Flash Attention、Ring Attention和Star Attention分別實(shí)現(xiàn)了高達(dá)9.2倍、4.2倍和1.6倍的速度提升。結(jié)合序列并行化與近似注意力機(jī)制，APB在保持任務(wù)性能的同時(shí)，顯著降低了計(jì)算量和通信開(kāi)銷。
• 高效的分布式計(jì)算：
  • 上下文分割：輸入序列被均勻分配到多個(gè)主機(jī)上，每個(gè)主機(jī)的本地上下文塊前添加一個(gè)錨點(diǎn)塊（Anchor Block），以保留對(duì)輸入序列初始部分的可見(jiàn)性。
  • 塊壓縮：在每個(gè)主機(jī)上，使用Locret的保留頭（Retaining Heads）對(duì)KV緩存進(jìn)行壓縮，從而減輕通信和計(jì)算負(fù)擔(dān)。
  • 通信機(jī)制：通過(guò)AllGather通信機(jī)制，將壓縮后的上下文塊發(fā)送到所有主機(jī)，并構(gòu)建傳遞塊（Passing Block），以傳遞前序主機(jī)的重要KV緩存單元。
  • 計(jì)算：每個(gè)主機(jī)結(jié)合錨點(diǎn)塊、傳遞塊和本地上下文塊進(jìn)行注意力計(jì)算。傳遞塊在注意力計(jì)算后被丟棄，不再參與后續(xù)計(jì)算。
• 適應(yīng)性強(qiáng)：APB支持多種模型和并行配置，能夠適應(yīng)不同的分布式設(shè)置和模型規(guī)模，具有良好的可擴(kuò)展性。通過(guò)調(diào)整錨點(diǎn)塊和傳遞塊的大小，APB可以在不同長(zhǎng)度的輸入序列上實(shí)現(xiàn)最佳性能。
• 保持任務(wù)性能：在長(zhǎng)上下文推理任務(wù)中，APB的速度更快，性能與全注意力計(jì)算（Full Attention）相當(dāng)，甚至在某些任務(wù)上表現(xiàn)更佳。通過(guò)查詢感知的上下文壓縮技術(shù)，APB能夠更加精確地識(shí)別和傳遞與查詢相關(guān)的上下文信息，從而保持或提升任務(wù)性能。

APB的技術(shù)原理

• 稀疏注意力機(jī)制：APB框架結(jié)合稀疏注意力機(jī)制，通過(guò)減少計(jì)算量來(lái)提升推理速度。實(shí)現(xiàn)稀疏注意力的方式包括：
  • 更小的Anchor block：與Star Attention相比，APB將Anchor block的大小縮小到上下文塊的1/4或1/8，從而降低了額外的計(jì)算開(kāi)銷。
  • Passing block：為了解決長(zhǎng)距離語(yǔ)義依賴問(wèn)題，APB通過(guò)構(gòu)建Passing block來(lái)傳遞重要信息。每個(gè)上下文塊被壓縮后，重要KV對(duì)通過(guò)通信傳遞到后續(xù)GPU上以構(gòu)建Passing block。
  • 查詢感知的上下文壓縮：APB在Anchor block的開(kāi)頭嵌入查詢，使上下文壓縮器能夠看到查詢內(nèi)容，從而更精準(zhǔn)地識(shí)別出與查詢相關(guān)的KV對(duì)，并通過(guò)通信機(jī)制傳遞給后續(xù)設(shè)備。
• 序列并行推理：APB框架采用序列并行的方式，將長(zhǎng)文本均勻分配到多個(gè)GPU上進(jìn)行并行處理，同時(shí)通過(guò)局部KV緩存壓縮及簡(jiǎn)化的跨GPU通信機(jī)制，解決了長(zhǎng)上下文中的遠(yuǎn)距離語(yǔ)義依賴問(wèn)題。

APB的項(xiàng)目地址

• Github倉(cāng)庫(kù)：https://github.com/thunlp/APB
• arXiv技術(shù)論文：https://arxiv.org/pdf/2502.12085

APB的應(yīng)用場(chǎng)景

• 長(zhǎng)文本推理：適用于長(zhǎng)文本生成、長(zhǎng)文本問(wèn)答等需要處理極長(zhǎng)輸入序列的應(yīng)用。
• 多Agent協(xié)作：適合多個(gè)Agent需要共同處理長(zhǎng)上下文信息的場(chǎng)景。
• 大規(guī)模模型服務(wù)：在分布式環(huán)境中高效處理長(zhǎng)上下文的模型服務(wù)。
• 知識(shí)圖譜構(gòu)建：在處理大量文本數(shù)據(jù)以提取和整合知識(shí)的任務(wù)中，APB框架通過(guò)高效的上下文壓縮和傳遞機(jī)制，能夠顯著提升知識(shí)圖譜構(gòu)建的效率。
• 實(shí)時(shí)交互系統(tǒng)：適用于需要快速處理用戶輸入并生成準(zhǔn)確回復(fù)的實(shí)時(shí)交互系統(tǒng)。APB框架通過(guò)高效的上下文壓縮和傳遞機(jī)制，顯著提升了實(shí)時(shí)交互系統(tǒng)的效率。