通过压缩上下文块跨GPU传递，加速长序列推理
APB: Accelerating Distributed Long-Context Inference by Passing Compressed Context Blocks across GPUs
- 采用多主机近似注意力机制，压缩关键KV对跨GPU通信
- 相比FlashAttn提速9.2倍，且不损失任务性能
- 适合需要高效处理超长文本的部署场景
长上下文推理对大语言模型应用至关重要，但其预填充阶段速度仍是主要瓶颈。现有方法如序列并行策略和近似注意力机制虽有改进，仍难以实现最优推理效率，限制了输入序列长度扩展和长上下文查询的实时处理。为此，我们提出APB框架，利用多主机近似注意力，在减少计算量的同时提升并行性，显著加速预填充阶段。APB在序列并行框架中引入关键值对的通信机制，实现更快推理速度且保持任务性能。通过集成定制FlashAttn内核与优化分发策略，APB支持多种模型与并行配置。实验表明，相比FlashAttn、RingAttn和StarAttn，APB分别实现最高9.2倍、4.2倍和1.6倍的加速，且无明显任务性能下降。代码已开源：https://github.com/thunlp/APB。
原文摘要 · Abstract (English)
While long-context inference is crucial for advancing large language model (LLM) applications, its prefill speed remains a significant bottleneck. Current approaches, including sequence parallelism strategies and compute reduction through approximate attention mechanisms, still fall short of delivering optimal inference efficiency. This hinders scaling the inputs to longer sequences and processing long-context queries in a timely manner. To address this, we introduce APB, an efficient long-context inference framework that leverages multi-host approximate attention to enhance prefill speed by reducing compute and enhancing parallelism simultaneously. APB introduces a communication mechanism for essential key-value pairs within a sequence parallelism framework, enabling a faster inference speed while maintaining task performance. We implement APB by incorporating a tailored FlashAttn kernel alongside optimized distribution strategies, supporting diverse models and parallelism configurations. APB achieves speedups of up to 9.2x, 4.2x, and 1.6x compared with FlashAttn, RingAttn, and StarAttn, respectively, without any observable task performance degradation. We provide the implementation and experiment code of APB in https://github.com/thunlp/APB.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。