DuoAttention 框架

DuoAttention 是由麻省理工学院 (MIT) 韩松团队于 2024 年提出的一种新型框架，旨在提高大型语言模型 (LLMs) 在处理长文本上下文时的推理效率。相关论文成果为「DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads」。

这个框架通过巧妙地区分「检索头 (Retrieval Heads)」和「流式头 (Streaming Heads)」2 种注意力头来优化模型的内存使用和计算速度。检索头专注于处理长距离依赖，需要完整的键值 (KV) 缓存，而流式头则关注最近 token 和注意力汇聚点，只需固定长度的 KV 缓存。这种设计显著减少了模型在解码和预填充过程中的内存使用和延迟，同时保持了模型处理长文本上下文的能力。

DuoAttention 通过为检索头应用全 KV 缓存，而为流式头应用轻量级、固定长度的 KV 缓存，实现了内存和计算资源的优化。这一改进不仅提升了模型解码速度和预填充效率，还减少了处理长文本时的延迟。对于多头注意力 (MHA) 模型最多可降低 2.55 倍，对于分组查询注意力 (GQA) 模型最多可降低 1.67 倍；同时，在解码速度方面，对于多头注意力 (MHA) 模型最多可提升 2.18 倍，对于分组查询注意力 (GQA) 模型最多可提升 1.50 倍；在预填充速度方面，对于多头注意力 (MHA) 模型最多可提升 1.73 倍，对于分组查询注意力 (GQA) 模型最多可提升 1.63 倍，而且与全注意力模式相比，精度损失极小。值得注意的是，结合量化技术，双注意力框架使得在单个 A100 GPU 上能够对具有 330 万文本长度的 Llama – 3 – 8B 模型进行解码。