分享自:

TurboRAG:基于预计算KV缓存的块文本检索增强生成加速方法

期刊:Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing

本文档属于类型a,是一份关于单一原创研究的学术论文报告。


关于TurboRAG:一种利用预计算KV缓存加速检索增强生成的新范式

一、 研究基本信息

本项研究的主要作者为宋硕璐(Songshuo Lu)、华天宇(Hua Wang)、容梓天(Yutian Rong)、陈志(Zhi Chen)和唐耀华(Yaohua Tang),他们均来自摩尔线程人工智能(Moore Threads AI)。该研究成果以题为《TurboRAG: Accelerating Retrieval-Augmented Generation with Precomputed KV Caches for Chunked Text》的论文形式,发表于计算语言学和自然语言处理领域的顶级会议——2025年自然语言处理经验方法会议(Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing, EMNLP 2025)。会议于2025年11月4日至9日举行,论文页码为6588至6601。

二、 学术背景与研究目标

本研究隶属于自然语言处理(NLP)领域,具体聚焦于检索增强生成(Retrieval-Augmented Generation, RAG)系统的推理加速技术。RAG通过将大型语言模型(LLM)与稠密检索器结合,使模型生成能够基于外部知识,已成为知识密集型NLP任务的事实标准方案。然而,现有的RAG系统采用“检索-拼接-预填充”的标准范式,即在线检索多个相关文档块,将其与用户查询拼接成一个长序列,然后由LLM进行一次性前向计算(预填充)以生成键值(Key-Value, KV)缓存,这导致了显著的首次令牌生成时间(Time-To-First-Token, TTFT)延迟。该延迟主要源于三个瓶颈:1) 冗余重复计算:频繁被检索的文档块需要在每次查询时重新编码;2) 二次方预填充成本:拼接多个文档块导致输入序列长度线性增长,而自注意力机制的计算成本随序列长度呈二次方增长;3) 受限的批处理大小:长上下文消耗大量GPU内存,限制了单设备批处理大小,从而影响吞吐量。

基于此背景,本研究旨在解决一个核心问题:能否将RAG的预填充过程转化为一种离线和在线混合的范式,通过预计算文档块的KV缓存并在线复用,来消除对检索内容在线预填充计算的需求?研究的目标是设计一种无需修改现有LLM架构和推理系统、即可显著降低TTFT并保持答案质量的新方法。

三、 研究详细工作流程

本研究的工作流程可分为离线和在线两个阶段,核心在于引入了“独立注意力”和“重排序位置编码”两项关键技术,以解决预计算缓存拼接时的语义一致性问题。

1. 离线预处理阶段: 此阶段针对固定的文档库进行一次性操作。 * 研究对象与处理:研究将整个文档库分割成固定大小的文本块(Chunk)。这些文本块是研究处理的基本单位。 * 关键步骤一:文档编码与索引。使用基于Transformer的模型(如BERT)将每个文档块编码为嵌入向量,并存储到向量数据库中,以支持后续高效的相似性检索。 * 关键步骤二:文档预填充与KV缓存计算。这是本研究的核心创新点之一。研究人员使用目标LLM(如Qwen2-7B)对每个独立的文档块进行预填充计算。与标准RAG不同,这里不是将多个块拼接后计算一个整体的KV缓存,而是为每一个单独的文档块计算其独立的KV缓存。计算时,采用标准的因果注意力掩码和从0开始的连续位置ID。计算完成后,这些块级别的KV缓存被持久化存储到数据库或磁盘中。这一步将原本需要在每次用户查询时在线进行的、针对长文档的计算,转移到了离线阶段。

2. 在线推理阶段: 此阶段响应用户的实时查询。 * 步骤一:查询编码。用户查询被相同的编码器(如BERT)转换为向量。 * 步骤二:检索。利用查询向量在向量数据库中进行相似性搜索,召回最相关的K个文档块。 * 步骤三:上下文KV缓存形成(在线拼接)。系统从存储中检索出与召回文档块对应的、预先计算好的KV缓存。此时面临核心挑战:简单地物理拼接这些独立计算的缓存会导致注意力掩码和位置索引不一致,从而严重影响模型输出质量。为此,研究提出了两项关键技术: * 独立注意力(Independent-Attention):在拼接后的整体上下文中,将不同文档块之间的交叉注意力权重强制设为零(即进行掩码)。如图2c所示,每个文档块内部的注意力保持正常,但块与块之间不可互相关注。这基于一个关键观察:在典型的RAG任务中,不同文档块之间的注意力权重本就非常稀疏,文档主要关注自身内部信息。 * 重排序位置编码(Reordered-RoPE):对于使用旋转位置编码(RoPE)的模型,其注意力计算仅依赖于查询和键之间的相对位置偏移,而非绝对位置索引。在预计算时,每个块的KV缓存是基于位置ID [0, 1, ..., L](L为块长度)计算的。如果直接拼接(称为“复合位置”),拼接后的位置ID序列会是 [0,...,L, 0,...,L, ...],这会破坏块内令牌与生成令牌之间的真实相对距离。因此,研究在在线推理时,在拼接缓存后,对整体的位置ID进行重排序,将其重新映射为一个从0开始的单调递增序列 [0,...,L, L+1,..., 2L, ...],并据此重新应用RoPE计算最终的位置感知键状态。这确保了相对位置关系的正确性。 * 步骤四:KV缓存预填充(在线)。LLM接收用户查询,并以上一步拼接并处理好的KV缓存作为上下文,进行轻量的预填充计算(主要计算查询的表示及其与上下文的注意力)。 * 步骤五:响应生成(在线)。预填充完成后,LLM开始自回归地生成答案并返回给用户。

3. 模型适应性训练: 为了确保LLM能够适应新的“独立注意力+重排序位置”的上下文格式,避免精度下降,研究对基础LLM进行了轻量级的监督微调(SFT)。训练数据混合了50%的文档问答数据和50%的通用任务数据(如代码、对话、推理)。训练使用相同的注意力掩码和位置ID设置(如图2b和2c所示),使模型学会在这种新的缓存布局下进行推理。

四、 主要实验结果

研究在多个标准基准测试上验证了TurboRAG的有效性,主要对比基线为标准的“朴素RAG”(Naïve RAG)系统。

1. 文档问答准确性: * RGB基准测试:在RGB双语基准测试上,评估了模型在不同噪声比例(即检索结果中不相关文档的比例)下的表现。关键发现如下: * 未经微调的TurboRAG-Reordered(使用重排序位置)表现稳健,平均准确率仅比朴素RAG下降4.2%(92.6 vs 96.8),即使在0.8的高噪声比下,下降也不超过6%。这表明该方法具备良好的开箱即用性。 * 未经微调的TurboRAG-Composite(使用复合位置)性能下降更明显,在高噪声下差距近20%,验证了保持正确相对位置关系的重要性。 * 经过微调后TurboRAG-ReorderedTurboRAG-Composite在RGB上的平均准确率与朴素RAG的差距均能控制在1%以内,与GPT-4o表现相当,证明了微调能有效消除因注意力机制改变带来的精度偏差。 * 在需要跨文档信息整合的任务上,TurboRAG-Reordered也达到了与朴素RAG相当的水平。 * LongBench基准测试:在长文本RAG基准LongBench上,TurboRAG同样表现出与朴素RAG可比拟的答案准确性,证明了该方法对长文本输入的有效性。 * 模型泛化性:在Llama-3.1-8B模型上的实验得到了与Qwen2-7B一致的结果,表明该方法可广泛应用于基于RoPE的LLM。

2. 推理速度(TTFT)性能: * 在LongBench多文档QA任务上,TurboRAG-Reordered相比朴素RAG,平均将TTFT降低了8.6倍,最高降低达9.4倍(例如,在Musique任务上,TTFT从1610ms降至171ms)。这主要归功于将长文本的KV缓存计算从在线转移到了离线。 * 即使在上下文较短(如1000个令牌)的情况下,TurboRAG也能实现约2.4倍的加速。 * 加速效果在不同规模的模型(Qwen2系列从1.5B到72B参数)上均保持一致,且随着批处理大小的增加,加速比会进一步提高(因为朴素RAG的GPU计算压力更大)。

3. 批处理扩展性与计算开销: * 在固定总上下文长度和查询长度的条件下,随着批处理大小从1增加到8,TurboRAG始终能保持约4倍以上的TTFT加速(即使考虑将KV缓存从CPU传输到GPU的开销)。如果缓存已预取至GPU内存,加速比可达16-33倍。 * 在计算量(TFLOPs)方面,TurboRAG相比朴素RAG减少了约98.5%的在线计算开销,这极大地释放了GPU计算资源,使得系统能够支持更大的批处理规模,从而提高总体吞吐量。

4. 通用能力回归测试: * 通过在OpenCompass基准上的多项主流任务(如MMLU、GSM8K、HumanEval等)测试表明,经过针对新注意力模式微调后的TurboRAG模型,其通用能力与朴素RAG模型相比没有出现显著退化,某些任务上甚至略有提升,证明了该方法的修改不会损害模型的原有能力。

五、 研究结论与价值

本研究提出并验证了TurboRAG,一种创新的混合离线-在线RAG推理范式。其核心结论是:通过预计算文档块的KV缓存,并在在线推理时利用“独立注意力”和“重排序RoPE”技术进行拼接,可以在不牺牲答案质量的前提下,显著降低RAG系统的首次令牌生成延迟(TTFT)和在线计算开销

该研究的价值体现在: * 科学价值:它挑战了传统RAG必须在线计算整个拼接上下文KV缓存的范式,证明了通过精心设计的注意力掩码和位置编码处理,可以实现缓存的高效复用而不损失模型性能。这为理解Transformer模型中跨文档注意力机制的必要性提供了新的实证依据(研究表明稀疏的跨块注意力可以被安全地屏蔽),也为LLM适应不同位置嵌入模式的灵活性提供了佐证。 * 应用价值:TurboRAG为对延迟敏感的大规模实时应用(如基于固定文档库的在线客服助手)和资源受限的边缘设备部署(如个人电脑、工作站)提供了高效的解决方案。它通过“以存储换计算”的策略,使得在计算资源有限的情况下也能实现流畅的交互体验。

六、 研究亮点

  1. 范式创新:首次系统性地将RAG推理中的在线预填充计算重构为离线和在线混合的范式,旨在从根本上消除对检索内容的在线预填充计算,而非仅仅优化在线计算本身。
  2. 关键技术简单有效:提出的“独立注意力”和“重排序位置编码”两项技术,概念清晰,实现相对简单,且能与现有大多数基于RoPE的LLM及其应用兼容,无需修改模型架构和底层推理系统。
  3. 显著的性能提升:实验证明了高达9.4倍的TTFT降低和98.5%的在线计算开销减少,效果非常显著。
  4. 保持精度与通用性:通过轻量级微调,成功解决了新缓存布局可能带来的精度损失问题,并在多个基准测试上证明了其答案质量与标准RAG相当,且不损害模型的通用能力。
  5. 明确的实用场景:研究明确指出了该方法在实时用户支持和资源受限设备两大场景下的巨大应用潜力。

七、 其他有价值内容

  • 局限性讨论:研究坦诚地讨论了TurboRAG的两个主要局限:1) 存储开销:需要为每个文档块存储KV缓存。研究通过成本分析指出,对于大多数实际场景(万到十万量级文档),存储成本是可接受的,且可与正交的KV缓存压缩技术结合进一步降低开销。2) 模型微调依赖:当前流程需要对基础模型进行微调,这限制了其在最新SOTA模型上的即时应用。作者将此列为未来的改进方向。
  • 与相关工作的对比:论文详细梳理了加速RAG的三条技术路线(减少检索文本量、加速解码本身、缓存复用),并将TurboRAG定位为缓存复用方向上的进展,其创新点在于通过位置重排序实现了顺序无关的缓存拼接。同时指出,TurboRAG与检索剪枝、缓存压缩、推测解码等技术是正交的,可以结合使用以获得更大收益。
  • 系统级部署考虑:论文在结论部分简要探讨了实际部署时可能面临的挑战,如动态知识库管理(版本化缓存、增量更新、在线预填充回退)、高性能KV缓存存储设计(利用NVMe SSD、LSM-tree索引、DMA优化)以及多节点设置等,为工程化落地提供了思路。
上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com