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

APB的主要功能

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

APB的技術原理

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

APB的應用場景

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