APB – 清華聯合騰訊等機構推出的分佈式長上下文推理框架

APB是什麼

APB（Accelerating Distributed Long-Context Inference by Passing Compressed Context Blocks across GPUs）是清華大學等機構聯合提出的分佈式長上下文推理框架。通過稀疏注意力機制和序列並行推理方式，有效解決了大模型處理長文本時的效率瓶頸。APB採用更小的Anchor block和Passing block，結合查詢感知的上下文壓縮技術，減少計算開銷的同時，精準傳遞關鍵信息，實現長距離語義依賴的高效處理。在128K文本上，APB推理速度比Flash Attention快約10倍，比英偉達的Star Attention快1.6倍，且性能優異。具備卓越的兼容性，能適應不同分佈式設定和模型大小。

APB的主要功能

• 加速長上下文推理：APB通過多主機近似注意力機制顯著提升推理速度，相比Flash Attention、Ring Attention和Star Attention分別實現了高達9.2倍、4.2倍和1.6倍的速度提升。通過序列並行化和近似注意力機制的結合，APB在保持任務性能的同時，大幅減少了計算量和通信開銷。
• 高效的分佈式計算：
  • 上下文分割：輸入序列被均勻分配到多個主機上，在每個主機的本地上下文塊前附加一個錨點塊（Anchor Block），保留對輸入序列初始部分的可見性。
  • 塊壓縮：在每個主機上，使用Locret的保留頭（Retaining Heads）對KV緩存進行壓縮，減少通信和計算開銷。
  • 通信機制：通過AllGather通信機制，將壓縮後的上下文塊發送到所有主機，並構建傳遞塊（Passing Block），以傳遞前序主機的重要KV緩存單元。
  • 計算：在每個主機上，結合錨點塊、傳遞塊和本地上下文塊進行注意力計算。傳遞塊在注意力計算後被丟棄，不參與後續計算。
• 適應性強：APB支持多種模型和並行配置，能適應不同的分佈式設置和模型大小，具有良好的可擴展性，通過調整錨點塊和傳遞塊的大小，APB可以在不同長度的輸入序列上實現最佳性能。
• 保持任務性能：在長上下文推理任務中，APB速度更快，在性能上與全注意力計算（Full Attention）相當，在某些任務上表現更好。通過查詢感知的上下文壓縮技術，APB能更精準地識別和傳遞與查詢相關的上下文信息，保持或提升任務性能。

APB的技術原理

• 稀疏注意力機制：APB框架整合了稀疏注意力機制，通過減少計算量來提升推理速度。通過以下方式實現稀疏注意力：
  • 更小的Anchor block：與Star Attention相比，APB將Anchor block的大小縮小到上下文塊的1/4或1/8，從而減少了額外的計算開銷。
  • Passing block：爲了解決長距離語義依賴問題，APB通過構建Passing block來傳遞重要信息。Passing block由前面設備上的重要KV對組成，每個上下文塊被壓縮後通信到後續GPU上構建Passing block。
  • 查詢感知的上下文壓縮：APB在Anchor block的開頭嵌入查詢，使上下文壓縮器能夠看到查詢的內容，更精準地識別出查詢相關的KV對，通過通信機制傳給後續設備。
• 序列並行推理：APB框架採用序列並行的方式，將長文本均勻分配到多個GPU上進行並行處理，同時通過局部KV緩存壓縮和精簡的跨GPU通信機制，解決了長上下文中的遠距離語義依賴問題。

APB的項目地址

• Github倉庫：https://github.com/thunlp/APB
• arXiv技術論文：https://arxiv.org/pdf/2502.12085

APB的應用場景

• 長文本推理：如長文本生成、長文本問答等，需要處理極長輸入序列的應用。
• 多Agent協作：多個Agent需要協同處理長上下文信息的場景。
• 大規模模型服務：需要在分佈式環境中高效處理長上下文的模型服務。
• 知識圖譜構建：知識圖譜構建任務需要處理大量的文本數據，提取和整合知識。APB框架通過高效的上下文壓縮和傳遞機制，能顯著提升知識圖譜構建的效率。
• 實時交互系統：實時交互系統需要快速處理用戶的輸入生成準確的回覆。APB框架通過高效的上下文壓縮和傳遞機制，能顯著提升實時交互系統的效率。