作者： AI 搜索引擎 来源： AI 搜索引擎

今天为大家分享北京大学和字节跳动联合发表的一篇关于加速RAG推理的论文：

RAGCache: Efficient Knowledge Caching for Retrieval-Augmented Generation

**论文地址：**https://arxiv.org/abs/2404.12457

1.论文概述

Retrieval-Augmented Generation (RAG) 是一种将检索机制与生成模型相结合的方法，通过首先检索与问题相关的文档或知识，然后生成模型基于检索结果生成答案或文本。最近的研究表明，RAG可以显著提高各种基准的生成质量。

然而，RAG的知识注入导致喂入大模型的序列信息增多，这引入了长序列生成的问题，导致较高的计算和存储开销。例如，初始请求包含100个token，检索到的文档总共包含1000个token。因此，增强请求的额外计算和内存成本比原始请求高10倍左右。资源需求的不断增加有效处理RAG请求提出了巨大挑战。

为了解决这个问题，本文设计了一个特定于RAG的多级动态缓存管理系统RAGCache 。本文发现了两个系统优化的机会。首先，相同文档在多个请求中重复出现，使得此类文档的LLM推理中间状态 (Key Value tensors) 可以共享。其次，一小部分文档占了大多数检索请求。因此，缓存这些频繁访问的文档的中间状态，可以减轻计算负担。

基于这个观察，RAG设计了一个知识树，它将检索到的文档的中间状态 (Key Value tensors) 适配到GPU和主机内存中。访问频率较高的文档缓存在快速的GPU内存中，访问频率较低的文档缓存在较慢的主机内存中。实验结果表明，与集成了Faiss的vLLM相比，RAGCache将TTFT时间降低了4倍，吞吐量提高了2.1倍。

2.核心方法

2.1 系统架构图

RAGCache架构如上图所示。

当接收到一个用户请求时，RAG Controller首先从外部数据库 (External Database) 检索相关的文档。然后，这些文档被转发到Cache Retriever 以查找匹配的Key Value tensors。如果缓存中没有Key Value tensors，RAGCache会指示LLM推理引擎 (LLM Inference Engine) 生成新的tokens。

相反，如果Key Value tensors可用，则带有Key Value tensors的请求将转发到LLM推理引擎，然后该引擎使用前缀缓存内核来生成标记。

生成第一个标记后，Key Value tensors被转发回RAG Controller，该控制器Key Value tensors并刷新缓存的状态。最后，生成的答案作为响应传递给用户。

2.2 缓存结构和替换策略

多个请求能够共享Key Value tensors的部分要求前缀相同，即必须保持相同的顺序。因此，为了在保持文档顺序的同时实现快速检索，RAGCache使用知识树来构建文档的键值张量，如下图所示。

该树将每个文档分配给一个节点，节点指的是文档的键值张量的内存地址。从根到特定节点的路径表示文档序列。RAGCache通过沿着这些路径的前缀匹配来检索张量。在前缀匹配过程中，如果后续文档不在子节点中，则立即终止遍历，并返回标识的文档序列。

替换策略 Prefix-aware Greedy-Dual-Size-Frequency (PGDSF)：利用知识树，RAGCache 必须确定每个节点在分层缓存中的位置。访问频率较高的节点最好存储在GPU内存中，以提高访问速度，而访问频率较低的节点则分配给速度较慢的主机内存或直接释放。

为了优化节点放置，RAGCache采用了PGDSF替换策略。PGDSF根据每个节点的访问频率、大小和访问成本计算优先级。此方法通过维护优先级最高的节点来利用有限的存储容量，其优先级定义如下：

时间越新优先级越高，访问频率越大优先级越高，平均成本越大优先级越高。

3.精彩段落

4.总结

RAGCache是专为RAG定制设计的新型多级动态缓存系统。RAGCache缓存了多个请求中检索到的文档的Key Value以最大限度地减少冗余计算。RAGCache的核心是一棵知识树，可确保缓存最关键的键值张量。

RAGCache还实现了一个全局RAG控制器，用于协调外部知识数据库和LLM推理引擎之间的交互。该控制器通过系统优化得到增强，包括缓存感知重新排序和动态推测流水线。

5.延展阅读

[1] SOSP'23：基于内存管理的LLM推理加速

[Efficient Memory Management for Large Language Model Serving with PagedAttention]

地址：https://dl.acm.org/doi/10.1145/3600006.3613165

👆 关注AI 搜索引擎 ，获取更多专业技术分享 ~

更多AI工具，参考Github-AiBard123，国内AiBard123