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通信机制，解决了长上下文中的远距离语义依赖问题。