❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发感兴趣，我会每日分享大模型与 AI 领域的开源项目和应用，提供运行实例和实用教程，帮助你快速上手AI技术！
🥦 AI 在线答疑 -> 智能检索历史文章和开源项目 -> 丰富的 AI 工具库 -> 每日更新 -> 尽在微信公众号 -> 搜一搜：蚝油菜花 🥦
🌟 "百万字文档秒级响应！清华用分布式Attention改写LLM游戏规则"
大家好，我是蚝油菜花。当ChatGPT还在为8K上下文挣扎时，中国团队已经让AI轻松驾驭百万字长文本！你是否正在经历：
- 👉 处理论文/法律合同时，模型推理慢如蜗牛
- 👉 尝试长上下文应用，显存瞬间爆炸报OOM
- 👉 用分布式方案却遭遇通信延迟拖后腿...
清华大学最新开源的 APB框架 ，正在颠覆长文本处理范式！这项黑科技：
- ✅ 碾压FlashAttention 10倍速度：128K文本推理快到飞起
- ✅ 智能上下文压缩：精准保留关键语义，性能反超全注意力
- ✅ 分布式无缝扩展：多GPU并行效率高达92%，百万字轻松拿捏
- ✅ 即插即用兼容性：适配主流大模型，无需重训即刻加速
某金融团队实测——200页财报分析从3小时缩至18分钟，关键信息提取准确率提升15%！想用开源代码复现这波操作？手把手教程现在开始！
🚀 快速阅读
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通过多主机近似注意力机制显著提升推理速度，相比Flash Attention、Ring Attention和Star Attention分别实现了高达9.2倍、4.2倍和1.6倍的速度提升。
- 高效的分布式计算：输入序列被均匀分配到多个主机上，在每个主机的本地上下文块前附加一个锚点块（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来传递重要信息。Passing block由前面设备上的重要KV对组成，每个上下文块被压缩后通信到后续GPU上构建Passing block。
- 查询感知的上下文压缩：APB在Anchor block的开头嵌入查询，使上下文压缩器能够看到查询的内容，更精准地识别出查询相关的KV对，通过通信机制传给后续设备。
- 序列并行推理：APB框架采用序列并行的方式，将长文本均匀分配到多个GPU上进行并行处理，同时通过局部KV缓存压缩和精简的跨GPU通信机制，解决了长上下文中的远距离语义依赖问题。
如何运行 APB
环境设置
conda create -n apb python=3.9
conda activate apb
pip install -r requirements.txt
pip install experiments/flash-attention-apb
pip install experiments/ring-flash-attention-main
下载 Retaining Heads
| BackBone Model | HF Repo |
|---|---|
| Llama-3.1-8B-instruct | huggingface.co/thunlp/APB-… |
| Qwen-2.5-14B-instruct | huggingface.co/thunlp/APB-… |
| Yi-34B-200K | huggingface.co/thunlp/APB-… |
| Llama-3-8B-1M-instruct | huggingface.co/thunlp/APB-… |
示例
我们提供了一个使用 APB 处理 NIAH-Simple-1 类查询的示例，使用 8 个 GPU。
首先，修改 example/llama.sh 中的模型路径、locret 路径和数字（NIAH 任务中的 needle）。
然后，运行以下命令。
bash example/llama.sh
输出预期如下（如果数字设置为 688435772345）：
Ground Truth: 688435772345
Prediction: 688435772345.
复现说明
请先设置环境，然后参考 experiments/README.md 获取详细信息。
资源
- GitHub 仓库：github.com/thunlp/APB
❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发感兴趣，我会每日分享大模型与 AI 领域的开源项目和应用，提供运行实例和实用教程，帮助你快速上手AI技术！
🥦 AI 在线答疑 -> 智能检索历史文章和开源项目 -> 丰富的 AI 工具库 -> 每日更新 -> 尽在微信公众号 -> 搜一搜：蚝油菜花 🥦
