APB - A distributed long-context reasoning framework jointly developed by Tsinghua University, Tencent, and other institutions.

What is APB?

APB (Accelerating Distributed Long-Context Inference by Passing Compressed Context Blocks across GPUs) is a distributed long-context inference framework jointly proposed by Tsinghua University and other institutions. Through sparse attention mechanisms and sequential parallel inference, it effectively addresses the efficiency bottleneck of large models processing long texts. APB employs smaller anchor blocks and passing blocks, combined with query-aware context compression technology, reducing computational overhead while accurately conveying key information, achieving efficient processing of long-distance semantic dependencies. On 128K text, APB inference speed is approximately 10 times faster than Flash Attention and 1.6 times faster than NVIDIA's Star Attention, with excellent performance. It also boasts excellent compatibility, adapting to different distributed setups and model sizes.

Main functions of APB

• Accelerate long context reasoningAPB significantly improves inference speed through a multi-host approximate attention mechanism, achieving speed improvements of up to 9.2x, 4.2x, and 1.6x compared to Flash Attention, Ring Attention, and Star Attention, respectively. By combining sequence parallelization and approximate attention, APB significantly reduces computational and communication overhead while maintaining task performance.
• High-efficiency distributed computing:
  • Context SegmentationThe input sequence is evenly distributed across multiple hosts, with an anchor block appended before the local context block on each host to preserve the visibility of the initial portion of the input sequence.
  • Block compressionOn each host, Locret's Retaining Headers are used to compress the KV cache, reducing communication and computational overhead.
  • Communication mechanismThe AllGather communication mechanism sends the compressed context block to all hosts and constructs a Passing Block to pass important key-value cache units from the preceding hosts.
  • calculateAttention is computed on each host by combining the anchor block, the transit block, and the local context block. The transit block is discarded after attention computation and does not participate in subsequent computations.
• Highly adaptableAPB supports multiple models and parallel configurations, can adapt to different distributed settings and model sizes, and has good scalability. By adjusting the size of the anchor block and the transit block, APB can achieve optimal performance on input sequences of different lengths.
• Maintain task performanceIn long-context reasoning tasks, APB is faster, performing comparably to Full Attention, and even better on some tasks. Through query-aware context compression, APB can more accurately identify and convey query-related contextual information, maintaining or improving task performance.

APB's technical principles

• Sparse attention mechanismThe APB framework integrates a sparse attention mechanism to improve inference speed by reducing computational cost. Sparse attention is implemented in the following ways:
  • Smaller Anchor BlockCompared to Star Attention, APB reduces the size of the Anchor block to 1/4 or 1/8 of the context block, thereby reducing additional computational overhead.
  • Passing blockTo address the problem of long-range semantic dependencies, APB constructs a Passing block to transmit important information. The Passing block consists of important key-value pairs from preceding devices, and each context block is compressed and communicated to subsequent GPUs to construct the Passing block.
  • Query-aware context compressionAPB embeds a query at the beginning of the Anchor block, enabling the context compressor to see the query content, more accurately identify the key-value pairs related to the query, and transmit them to subsequent devices through a communication mechanism.
• Sequence Parallel InferenceThe APB framework uses a sequence parallel approach to distribute long texts evenly across multiple GPUs for parallel processing. It also solves the problem of long-distance semantic dependencies in long contexts through local key-value caching compression and a simplified cross-GPU communication mechanism.

APB project address

• Github repository:https://github.com/thunlp/APB
• arXiv technical paper:https://arxiv.org/pdf/2502.12085

APB application scenarios

• Long text reasoningApplications that require handling extremely long input sequences, such as long text generation and long text question answering.
• Multi-Agent Collaboration: Scenarios where multiple agents need to collaboratively process long context information.
• Large-scale model servicesModel services that need to efficiently handle long contexts in a distributed environment.
• Knowledge Graph ConstructionKnowledge graph construction requires processing large amounts of text data to extract and integrate knowledge. The APB framework can significantly improve the efficiency of knowledge graph construction through its efficient context compression and transmission mechanisms.
• Real-time interactive systemReal-time interactive systems need to quickly process user input and generate accurate responses. The APB framework can significantly improve the efficiency of real-time interactive systems through its efficient context compression and transmission mechanisms.