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 institutions such as Tsinghua University. Through a sparse attention mechanism and sequence-parallel inference approach, it effectively addresses the efficiency bottleneck when large models process long texts. APB uses smaller Anchor blocks and Passing blocks, combined with query-aware context compression technology, to reduce computational overhead while precisely transmitting key information, enabling efficient handling of long-distance semantic dependencies. On 128K text, APB inference speed is about 10 times faster than Flash Attention and 1.6 times faster than NVIDIA's Star Attention, with excellent performance. It has exceptional compatibility and can adapt to different distributed settings and model sizes.
Main Features of APB
- Accelerated Long-Context Inference: APB significantly improves inference speed through a multi-host approximate attention mechanism, achieving speedups of up to 9.2x, 4.2x, and 1.6x compared with Flash Attention, Ring Attention, and Star Attention, respectively. By combining sequence parallelization and an approximate attention mechanism, APB greatly reduces computation and communication overhead while maintaining task performance.
- Efficient Distributed Computing：
- Context Segmentation: The input sequence is evenly distributed across multiple hosts, and an Anchor Block is appended before each host's local context block, preserving visibility of the initial part of the input sequence.
- Block Compression: On each host, the KV cache is compressed using Locret's Retaining Heads to reduce communication and computation overhead.
- Communication Mechanism: Through the AllGather communication mechanism, the compressed context blocks are sent to all hosts, and a Passing Block is constructed to transmit the important KV cache units from preceding hosts.
- Computation: On each host, attention computation is performed by combining the Anchor Block, Passing Block, and local context block. The Passing Block is discarded after the attention computation and does not participate in subsequent computation.
- Strong Adaptability: APB supports multiple models and parallel configurations, can adapt to different distributed settings and model sizes, and has good scalability. By adjusting the sizes of the Anchor Block and Passing Block, APB can achieve optimal performance on input sequences of different lengths.
- Maintains Task Performance: In long-context inference tasks, APB is faster and performs comparably to Full Attention, and performs better on some tasks. Through query-aware context compression technology, APB can more accurately identify and transmit context information relevant to the query, maintaining or improving task performance.
Technical Principles of APB
- Sparse Attention Mechanism: The APB framework integrates a sparse attention mechanism to improve inference speed by reducing computation. Sparse attention is implemented in the following ways:
- Smaller Anchor block: Compared with 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 block: To address the problem of long-distance semantic dependencies, APB constructs a Passing block to transmit important information. The Passing block consists of important KV pairs on previous devices, and each context block is compressed and communicated to subsequent GPUs to construct the Passing block.
- Query-Aware Context Compression: APB embeds the query at the beginning of the Anchor block, enabling the context compressor to see the content of the query, more accurately identify query-related KV pairs, and pass them to subsequent devices through a communication mechanism.
- Sequence-parallel inference: The APB framework adopts a sequence-parallel approach, evenly distributing long text across multiple GPUs for parallel processing. At the same time, through local KV cache compression and a streamlined cross-GPU communication mechanism, it solves the problem of long-distance semantic dependencies in long contexts.
APB project address
- Github repository：https://github.com/thunlp/APB
- arXiv technical paper：https://arxiv.org/pdf/2502.12085
APB application scenarios
- Long-text inference: Such as long-text generation, long-text question answering, and other applications that need to process extremely long input sequences.
- Multi-Agent collaboration: Scenarios where multiple Agents need to collaboratively process long-context information.
- Large-scale model services: Model services that need to efficiently process long contexts in a distributed environment.
- Knowledge graph construction: Knowledge graph construction tasks need to process large amounts of text data to extract and integrate knowledge. Through an efficient context compression and transfer mechanism, the APB framework can significantly improve the efficiency of knowledge graph construction.
- Real-time interactive systems: Real-time interactive systems need to quickly process user input and generate accurate responses. Through an efficient context compression and transfer mechanism, the APB framework can significantly improve the efficiency of real-time interactive systems.