APB

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

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來傳遞重要信息。每個上下文塊被壓縮后，重要KV對通過通信傳遞到后續GPU上以構建Passing block。
- 查詢感知的上下文壓縮：APB在Anchor block的開頭嵌入查詢，使上下文壓縮器能夠看到查詢內容，從而更精準地識別出與查詢相關的KV對，并通過通信機制傳遞給后續設備。
序列并行推理：APB框架采用序列并行的方式，將長文本均勻分配到多個GPU上進行并行處理，同時通過局部KV緩存壓縮及簡化的跨GPU通信機制，解決了長上下文中的遠距離語義依賴問題。