APB-清华联合腾讯等机构推出的分布式长上下文推理框架

发布于:2025-03-19 ⋅ 阅读:(10) ⋅ 点赞:(0)

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框架通过高效的上下文压缩和传递机制,能显著提升实时交互系统的效率。