当前位置：首页 > news >正文

DUOATTENTION：结合检索与流式注意力机制的高效长上下文大语言模型推理方法

news 2025/11/6 8:00:57

摘要

部署长上下文的大语言模型（LLMs）是必要的，但也带来了显著的计算与内存挑战。将所有注意力头的 Key 和 Value（KV）状态进行缓存会占用大量内存。现有的 KV 缓存裁剪方法要么破坏了 LLM 的长上下文能力，要么仅带来有限的效率提升。本文发现，只有部分注意力头（即“检索头”）在处理长上下文时至关重要，需要对所有 token 执行完整注意力；而其余的注意力头主要关注最近的 token 和注意力汇聚点（attention sinks），我们称之为“流式头”，它们并不需要完整注意力机制。

基于这一洞察，我们提出 DuoAttention 框架：该框架仅对检索头使用完整的 KV 缓存，而对流式头则采用轻量级、固定长度的 KV 缓存，从而在不影响模型长上下文能力的前提下，降低 LLM 的解码和预填充时的内存占用与延迟。DuoAttention 借助轻量级、基于优化的算法，并利用合成数据准确识别检索头。我们的方法在保持精度几乎无损的前提下，将长上下文推理的内存使用量最多降低至原来的 2.55×（用于 MHA 模型）和 1.67×（用于 GQA 模型），同时解码速度最多提升至 2.18× 和 1.50×，预填充速度最多提升至 1.73× 和 1.63×。值得注意的是，结合量化技术后，DuoAttention 使得 Llama-3-8B 在单个 A100 GPU 上支持 330 万 token 的上下文长度推理成为可能。代码地址详见文末链接。