RAGCache: Efficient Knowledge Caching for Retrieval-Augmented Generation

发表时间: 2024-04 · arXiv:2404.12457 (TOCS 2025)

原文: https://arxiv.org/abs/2404.12457

作者/机构:Chao Jin1, Zili Zhang1, Xuanlin Jiang1, Fangyue Liu1, Xin Liu2, Xuanzhe Liu1, Xin Jin1 (1 北京大学, 2 字节跳动)

速读

一句话结论 提出了一种专为 RAG 定制的多级动态缓存系统 RAGCache,通过在 GPU 和主机内存中跨请求缓存检索文档的 KV cache,并将检索与推理流水线化,将首 token 延迟最高降低了 4 倍,吞吐量最高提升了 2.1 倍。

要解决什么问题 RAG 系统在工作时,会将检索到的外部文档注入到原始请求中,导致输入序列极长(例如请求只有 100 个 token,但注入文档可能长达数千个 token)。这使得大模型推理的 prefill(预填充)阶段需要消耗大量的计算资源和显存来生成这些长序列的 KV cache,成为整个系统的性能瓶颈。现有的 LLM 推理系统或缓存框架虽然支持 KV cache 复用,但它们仅将状态缓存在容量有限的 GPU 显存中,无法容纳 RAG 场景下海量的文档状态。此外,RAG 的检索步骤(在 CPU 上运行)和生成步骤(在 GPU 上运行)通常是串行执行的,导致检索期间 GPU 资源闲置,进一步拉长了端到端延迟。作者观察到 RAG 检索具有高度倾斜的分布特征(少数文档被频繁访问),且大模型对文档的拼接顺序极其敏感,这为设计专用的多级缓存提供了空间。

怎么做的 RAGCache 的核心思路是将高频检索文档的 KV cache 组织成树状结构,并在 GPU 和主机内存之间进行多级动态调度,从而避免跨请求的重复计算。系统主要由四个关键设计构成。第一是知识树(Knowledge Tree),它以文档 ID 为节点构建前缀树。因为 LLM 的注意力机制对位置敏感(文档 A 在 B 前面,和 B 在 A 前面的 KV cache 是不同的),知识树通过严格的前缀匹配来保证文档顺序的正确性,并允许不同请求共享相同的前缀路径。第二是前缀感知替换策略(PGDSF),用于在显存和内存间驱逐或交换节点。它综合考虑了节点的访问频率、大小、最近访问时间以及重算成本,其优先级定义为:

$$Priority = Clock + \frac{Frequency \times Cost}{Size}$$


为了适应 RAG 的前缀特性,单位大小的重算成本被设计为平摊到所有未缓存该文档的请求上:

$$\frac{Cost}{Size} = \frac{1}{m} \sum_{i=1}^{m} \frac{Cost_i}{NewSize_i}$$
优先级最低的节点会先从 GPU 显存被交换到主机内存,若主机内存也满则被彻底释放。第三是缓存感知重排(Cache-aware Reordering),当请求并发较高时,调度器会优先处理那些能最大化缓存命中率的请求,其调度优先级定义为已缓存长度与需计算长度的比值:
$$OrderPriority = \frac{CachedLength}{ComputationLength}$$
第四是动态推测流水线(Dynamic Speculative Pipelining)。为了打破检索和生成的串行阻塞,系统会在向量数据库检索的中间阶段,将初步的 Top-k 文档提前送入 GPU 进行推测性生成。如果最终检索结果与初步结果一致,则直接返回生成的 token;如果不一致,则终止当前推测并重新生成。为了防止高负载下推测错误浪费算力,系统仅在 GPU 处于空闲或请求池未满时才启动该流水线。

效果如何 实验在 AWS EC2 实例(单卡 A10G 24GB 显存配合 192GB 主机内存)和双卡 H800 80GB 环境下进行,使用了 LLaMA2-7B、Mistral-7B 以及更大规模的 Mixtral-8x7B 和 LLaMA2-70B 模型。外部知识库基于维基百科构建,包含约 30 万篇文档,并在 MMLU 和 Natural Questions 问答数据集上生成模拟请求。对比基线有两个点名系统:代表底层显存优化路线的 vLLM(采用 PagedAttention 减少显存碎片),以及代表 GPU 内 KV cache 复用路线的 SGLang(采用 LRU 策略)。量化结果显示,在 MMLU 数据集上,RAGCache 相比 vLLM 将首 token 延迟(TTFT)降低了 1.2 到 4 倍,吞吐量提升了 1.3 到 2.1 倍;相比 SGLang,TTFT 降低了 1.1 到 3.5 倍,吞吐量提升了 1.2 到 1.8 倍。在代价与局限性方面,系统依赖主机内存作为二级缓存,当发生缓存命中但数据在主机内存时,需要通过 PCIe 总线将 KV cache 传输到 GPU,这会产生一定的传输延迟,但实验证明该延迟仍比完全重新计算快 3.9 倍。此外,推测流水线在检索准确率较低的早期阶段可能会产生无效计算,因此必须依赖动态负载监控来决定是否开启。

主要贡献

核心问题:检索增强生成(RAG)通过整合大型语言模型(LLMs)和外部知识库的优势,在各种自然语言处理任务中展现了显著的改进。然而,在RAG工作流中,由于需要将检索到的文档(即外部知识)注入到原始请求中,引入了长序列生成问题,这导致了极高的计算和内存开销。具体而言,注入外部知识后的增强请求在计算和内存需求上可能比原始请求高出 $10\times$ 以上,这为高效处理RAG请求的系统扩展带来了巨大挑战。

研究目标:分析当前RAG系统的性能瓶颈(知识注入导致的长序列)和优化机会(缓存知识的中间状态),并设计一个专门为RAG定制的缓存系统,以降低延迟并提高吞吐量。

创新点
1. 系统特征分析:对RAG进行了详细的系统级特征分析,揭示了性能瓶颈和优化机会。
2. RAGCache系统:提出了RAGCache,这是首个通过缓存外部知识的中间状态(Key-Value Tensors)并在多个查询间共享这些状态以减少冗余计算的RAG系统。
3. 知识树与PGDSF替换策略:设计了知识树结构来组织GPU和主机内存层次结构中的中间状态,并提出了一种前缀感知(Prefix-aware)的Greedy-Dual-Size-Frequency (PGDSF) 缓存替换策略,该策略利用RAG的特性(文档顺序、大小、频率和最近访问时间)来最小化缓存未命中率。
4. 动态推测流水线(Dynamic Speculative Pipelining):提出了一种动态推测流水线方法,动态重叠检索和推理步骤的计算,以最小化端到端延迟。
5. 系统实现与卓越性能:实现了一个RAGCache原型系统。评估表明,与集成了Faiss的vLLM相比,RAGCache将首个Token生成时间(TTFT)降低了高达 $4\times$,吞吐量提高了高达 $2.1\times$。

图1 RAG工作流程
图1 RAG工作流程

背景知识与关键观察

RAG的两步工作流:RAG结合了LLMs的深度上下文理解和知识库检索的精确性。其工作流分为检索和生成两步。在离线阶段,外部文档通过嵌入模型转化为高维向量并建立索引。当接收到用户请求时,RAG首先在向量数据库中进行向量相似度搜索,检索出最匹配的文档。接着,RAG将这些文档内容与原请求结合,生成增强请求并输入给LLM以生成响应。检索步骤主要在CPU上执行,而生成步骤在GPU上执行。

LLM生成的性能瓶颈:LLM推理分为预填充(Prefill)和解码(Decoding)两个阶段。预填充阶段需要计算整个输入序列的键值(Key-Value)张量,非常耗时。由于RAG增强了请求长度(文档平均长度远超请求长度),生成步骤(尤其是预填充阶段)的耗时随着序列长度急剧增加。当序列长度超过4000个Token时,推理时间达到一秒,成为系统的主要性能瓶颈。
图2 不同输入长度下的推理时间
图3 Token数量分布

缓存优化的机会:通过缓存先前检索文档的键值张量,可以显著减少后续相同文档请求的预填充延迟。实验表明,缓存前缀后的预填充延迟比全量预填充延迟低高达 $11.5\times$。即使考虑将KV Cache从主机内存传输到GPU内存的开销,缓存命中延迟依然比全量计算低高达 $3.9\times$。
图4 预填充延迟特征分析

偏斜的检索模式:对四个代表性问答数据集(MMLU、Google Natural Questions、HotpotQA、TriviaQA)的分析显示,文档检索模式呈现高度偏斜。少部分文档占据了绝大多数的检索请求(例如在MMLU中,前 $3\%$ 的文档被 $60\%$ 的请求引用)。这种低未命中率的特性使得缓存高频访问文档具有极大的潜力。
图5 不同数据集上的检索模式
图6 不同设置下的检索模式

方法细节

系统架构概览:RAGCache是一个为RAG量身定制的新型多级动态缓存系统。当请求到达时,RAG控制器首先从外部知识库检索相关文档。然后,这些文档被转发给缓存检索器(Cache Retriever),以在内存缓存中定位匹配的键值张量。如果缓存中不存在该张量,RAGCache会指示LLM推理引擎生成新的Token;如果存在,则将请求和对应的键值张量一起转发给LLM推理引擎,推理引擎利用前缀缓存内核进行Token生成。在生成首个Token后,键值张量被传回RAG控制器,控制器缓存这些被访问文档的张量并刷新缓存状态。最后,生成的答案被返回给用户。
图7 RAGCache概览

缓存检索器设计:缓存检索器利用知识树(Knowledge Tree)来高效定位内存中存储的文档键值张量。这棵树是基于文档ID构建的前缀树,完全契合LLM对文档顺序的位置敏感性。树中的每条路径代表一个请求引用的特定文档序列,每个节点保存被引用文档的键值张量。不同路径可以共享相同的节点,这表明不同请求之间可以共享文档。这种结构使得检索器能够迅速按照指定顺序访问文档的键值张量。

RAG控制器功能:RAG控制器负责协调交互,并包含几项专门针对RAG的系统优化。控制器采用了前缀感知的Greedy Dual-Size Frequency (PGDSF) 策略来最小化缓存未命中率。PGDSF基于频率、键值张量大小、最后访问时间以及前缀感知的重计算成本来计算优先级。缓存驱逐由该优先级决定,从而确保保留最有价值的张量。此外,控制器还实现了缓存感知重排序(Cache-aware reordering)以提高缓存命中率并防止缓存颠簸,同时保证请求公平性以缓解饥饿问题。动态推测流水线(Dynamic speculative pipelining)则被设计用来重叠知识检索和LLM推理,通过利用检索结果的中间生成状态提前启动LLM推理,从而最小化延迟。

缓存结构与顺序敏感性:与缓存独立对象的传统系统不同,RAGCache缓存的是对引用顺序敏感的检索文档键值张量。例如,考虑两个文档序列 $[D_1, D_3]$(键值张量为 $KV$)和 $[D_2, D_3]$(键值张量为 $KV'$)。尽管 $KV[1]$ 和 $KV'[1]$ 都属于文档 $D_3$,但它们的值是不同的。这种差异是因为给定Token的键值张量是基于前面的Token生成的,这突显了键值张量的顺序依赖性。

知识树的组织方式:为了在保持文档顺序的同时实现快速检索,RAGCache利用知识树来结构化文档的键值张量。该树将每个文档分配给一个节点,节点指向该文档键值张量的内存地址。遵循vLLM【索引编号1,Efficient memory management for large language model serving with pagedattention+2023+SOSP】的设计,RAGCache将键值张量存储在非连续的内存块中以实现KV Cache重用。根节点 $S$ 表示共享的系统提示词(System Prompt)。从根节点到特定节点的路径代表一个文档序列。
图8 知识树

前缀匹配与路径共享:知识树的设计天然允许RAGCache通过树中重叠的路径同时服务多个请求。RAGCache通过沿着这些路径进行前缀匹配来检索张量。在前缀匹配过程中,如果后续文档未能在子节点中找到,遍历会立即终止,并返回已识别的文档序列。这种方法确保了高效性,其时间复杂度为 $O(h)$,其中 $h$ 代表树的高度。

PGDSF缓存替换策略:随着知识树的运行,RAGCache必须决定每个节点在分层缓存中的放置位置。理想情况下,访问频繁的节点存储在GPU内存中以获得更快的访问速度,而访问较少的节点则分配到较慢的主机内存中或直接释放。为了优化节点放置,RAGCache采用了一种前缀感知的Greedy-Dual-Size-Frequency (PGDSF) 替换策略,该策略基于经典的GDSF策略【索引编号2,Improving WWW proxies performance with greedy-dual-size-frequency caching policy+1998】。与忽略文档大小变化的LRU等传统策略不同,PGDSF根据节点的访问频率、大小和访问成本来评估每个节点。该方法通过保留最有益的节点来利用有限的存储容量,节点的优先级定义如下:

$$Priority = Clock + \frac{Frequency \times Cost}{Size}$$

逻辑时钟机制:优先级较低的节点会被优先驱逐。$Clock$ 用于跟踪节点的访问新近度。我们在RAG控制器中为GPU和主机内存分别维护两个独立的逻辑时钟,以适应缓存层次结构。每个时钟从零开始,并在每次驱逐时更新。当一个文档被检索时,其节点的时钟会被设置,优先级也会相应调整。时钟较旧(表明最近较少使用)的节点优先级较低。设 $E$ 为一次驱逐操作中被驱逐节点的集合,时钟更新公式如下:

$$Clock = \max_{n \in E} Priority(n)$$

频率与大小的定义:$Frequency$ 表示一个时间窗口内文档的总检索次数。在系统启动或缓存清理时,该计数重置为零。优先级与频率成正比,因此访问越频繁的文档优先级越高。$Size$ 反映了文档分词后的Token数量,直接影响其键值张量所需的内存。$Cost$ 定义为计算文档键值张量所需的时间,它随GPU计算能力、文档大小以及前置文档序列的变化而变化。

前缀感知的成本估算:PGDSF在成本估算和节点放置两个方面实现了针对RAG系统的前缀感知。与GDSF中成本明确(如Web缓存中的对象大小)不同,RAG的成本涉及复杂的LLM生成动态。例如,图9展示了同一请求 $[S, D_1, D_2, Q]$ 产生的不同成本。为了估算 $D_2$ 的成本,直接使用缓存了 $[S, D_1]$ 或仅缓存了 $S$ 时的成本是不准确的。此外,后一种情况的成本还包含了计算 $D_1$ 和 $Q$ 键值张量的时间。PGDSF通过以下方式替换优先级公式中的 $\frac{Cost}{Size}$ 来解决这个问题:

$$\frac{Cost}{Size} = \frac{1}{m} \sum_{i=1}^{m} \frac{Cost_i}{NewSize_i}$$


其中 $m$ 是访问该文档但未缓存该文档的请求数量。$Cost_i / NewSize_i$ 表示第 $i$ 个请求中每个未缓存Token的计算时间。这种估算通过将成本分摊到所有未缓存的Token上,内在考虑了文档大小。至于 $Cost_i$,RAGCache离线分析了在不同缓存和非缓存Token长度下的LLM预填充时间,并使用双线性插值来估算给定请求的成本。文档检索会触发知识树中节点频率、成本估算和时钟的更新,或者为先前未缓存的文档初始化新节点。
图9 PGDSF中的成本估算

分层节点放置与驱逐:PGDSF协调知识树中的节点放置,知识树被划分为GPU、主机和空闲段。GPU内存中的节点作为主机内存中节点的父节点,建立了一种层次结构。RAGCache动态管理这些段之间的节点驱逐以提高效率。具体而言,当GPU内存已满时,RAGCache将叶子节点中优先级最低的节点交换到主机内存中。RAGCache对主机内存超额订阅应用类似的过程。这种驱逐策略维持了树的分层分区,这对于对齐内存层次结构和LLM生成中的前缀敏感性至关重要。一个节点依赖其父节点进行键值张量计算,这强调了优先放置父节点以实现快速检索的必要性。

更新与驱逐算法:算法1概述了在知识树GPU内存中更新和驱逐节点的操作。$T(\alpha, \beta)$ 表示拥有 $\alpha$ 个缓存Token和 $\beta$ 个非缓存Token的请求的估算计算时间。当请求检索到文档时,如果文档未被缓存,RAGCache会使用双线性插值更新成本;EVICT_IN_GPU 会从GPU内存中驱逐节点以容纳新请求,并根据公式更新时钟。如果父节点在驱逐后变成叶子节点,它将被添加到候选集合中。

function UPDATE_NODE_IN_GPU(node, is_cached, alpha, beta):
    # alpha and beta are cached and non-cached sizes of the request
    node.Frequency = node.Frequency + 1
    if is_cached is false:
        # Bilinear interpolation to estimate the cost
        Find alpha1 < alpha < alphah and beta1 < beta < betah from the profiler
        T_beta1 = T(alpha1, beta1) + (alpha - alpha1)/(alphah - alpha1) * [T(alphah, beta1) - T(alpha1, beta1)]
        T_betah = T(alpha1, betah) + (alpha - alpha1)/(alphah - alpha1) * [T(alphah, betah) - T(alpha1, betah)]
        T(alpha, beta) = T_beta1 + (beta - beta1)/(betah - beta1) * (T_betah - T_beta1)
        node.TotalCost = node.TotalCost + T(alpha, beta)
        node.numComputed = node.numComputed + 1
        node.AvgCost = node.TotalCost / node.numComputed
    node.Priority = Clock + node.AvgCost * node.Frequency

function EVICT_IN_GPU(required_size):
    E = {} # Evicted nodes in GPU
    S = {n in GPU | n.Children not in GPU} # Leaf nodes in GPU
    while sum(n.Size for n in E) < required_size:
        n = argmin_{n in S} n.Priority
        E = E union {n}
        Clock = max(Clock, n.Priority)
        if n.Parent.Children not in GPU:
            S = S union {n.Parent}

仅换出一次策略(Swap out only once):GPU通过PCIe总线连接到主机内存,其带宽通常远低于GPU HBM。为了最小化GPU和主机内存之间的数据传输,RAGCache采用了仅换出一次策略。节点的键值张量仅在首次驱逐时被换出到主机内存。主机内存保留这些键值张量,直到该节点从整个缓存中被驱逐。对于后续在GPU内存中的驱逐,RAGCache直接释放节点,实现零数据拷贝。考虑到主机内存的容量比GPU内存大一到两个数量级,在主机内存中保留一份键值张量的副本是可以接受的。

缓存感知重排序的动机:缓存命中率对RAGCache的缓存效率至关重要,但用户请求不可预测的到达模式会导致严重的缓存颠簸。引用相同文档的请求可能不会一起发出,从而影响缓存效率。例如,假设请求 $\{Q_i, i\%2==0\}$ 和 $\{Q_i, i\%2==1\}$ 分别指向文档 $D_1$ 和 $D_2$,且缓存容量为一个文档。序列 $\{Q_1, Q_2, Q_3...\}$ 会导致 $D_1$ 和 $D_2$ 的KV Cache频繁交换,缓存命中率为零。相反,将请求重新排序为 $\{Q_1, Q_3, Q_5, Q_2, Q_4, Q_6, Q_7, ...\}$ 可以优化缓存利用率,将命中率提高到 $66\%$。这说明了策略性的请求排序如何减轻缓存波动并提高缓存效率。

重排序的两个洞察:在介绍缓存感知重排序算法之前,我们先考虑两个场景。假设重计算成本与重计算长度成正比。第一个场景(图10a)考虑具有相同重计算需求但缓存上下文长度不同的请求,缓存限制为4。如果初始顺序为 $\{Q_1, Q_2\}$,系统必须清除 $Q_2$ 的缓存空间以容纳 $Q_1$ 的计算,然后重新为 $Q_1$ 分配内存。这有效利用了 $Q_1$ 的缓存但丢弃了 $Q_2$ 的,导致总计算成本为 $2+1+2=5$。相反,排序为 $Q_2, Q_1$ 会利用 $Q_2$ 的缓存但丢弃 $Q_1$ 的,计算增加到 $2+2+2=6$。因此,缓存感知重排序主张优先处理具有更长缓存上下文的请求。第二个场景(图10b)检查具有相同缓存上下文长度但重计算需求不同的请求,缓存容量为5。对于序列 $\{Q_1, Q_2\}$,系统必须清除 $Q_2$ 的缓存以为 $Q_1$ 分配空间。这需要完全重新计算 $Q_2$,总成本为 $2+2+1=5$。相反,序列 $\{Q_2, Q_1\}$ 允许直接计算 $Q_2$,将总计算减少到 $2+1=3$。因此,优先处理重计算段较短的请求是有益的。
图10 缓存感知重排序

重排序算法实现:基于这些洞察,我们引入了旨在提高缓存效率的缓存感知重排序算法。RAGCache使用优先级队列来管理传入的请求,根据它们对缓存性能的影响进行优先级排序。具体而言,请求的处理顺序基于以下优先级指标:

$$OrderPriority = \frac{CachedLength}{ComputationLength}$$


该公式优先考虑那些可能提高缓存效率的请求——即相对于计算需求,具有较大部分已缓存内容的请求。通过采用这种缓存感知重排序,RAGCache提高了缓存命中率并减少了总计算时间。为了避免饥饿,RAGCache为每个请求设置了一个窗口,确保所有请求在不晚于窗口大小的时间内得到处理。

动态推测流水线的动机:LLM生成是RAG系统中的关键性能瓶颈。然而,如果向量数据库规模变大或检索需要更高精度,检索步骤可能会产生大量延迟。为了减轻检索延迟的影响,RAGCache采用动态推测流水线来重叠知识检索和LLM推理。该技术背后的核心洞察是,向量搜索可能在检索步骤的早期就产生最终结果,LLM可以利用这些结果提前进行推测性生成。

推测流水线工作机制:具体而言,向量搜索维护一个Top-$k$ 候选文档队列,这些文档按与请求的相似度排序。在检索过程中,队列中的Top-$k$文档不断更新。然而,最终的Top-$k$文档可能在检索步骤的早期出现【索引编号3,Improving approximate nearest neighbor search through learned adaptive early termination+2020+SIGMOD】。基于这一观察,RAGCache将请求的检索过程分为多个阶段。在每个阶段,RAGCache触发向量数据库将候选文档发送给LLM引擎进行推测生成。然后,LLM引擎启动新的推测生成,如果接收到的文档与之前的不同,则终止之前的生成;如果相同,LLM引擎继续处理之前的生成。当最终的Top-$k$文档产生时,RAGCache将最终结果发送给LLM引擎。此时,如果最新推测生成的结果与最终Top-$k$文档匹配,LLM引擎直接将结果返回给用户;否则,LLM引擎执行重新生成。
图11 推测流水线

推测流水线的动态控制:推测流水线允许RAGCache重叠检索和生成步骤,从而降低RAG系统的端到端延迟。然而,它可能会引入额外的LLM计算,因为某些推测生成是不正确的,这在高系统负载下可能导致性能下降。为了解决这个问题,RAGCache根据系统负载动态启用推测流水线。我们首先通过简化分析(假设向量搜索和LLM一次只服务一个请求,使用G/G/1队列模型【索引编号4,Fundamentals of queueing theory+2018】)来证明如何最小化端到端延迟,同时控制系统负载。
图12 简化分析中的最优推测流水线策略

定理5.1与证明:设 $d$ 为推测时间间隔,$T$ 为LLM在G/G/1请求池中的服务时间分布,满足所有服务事件 $T \geq d$。最优策略是:如果在产生不同文档序列的阶段结束后请求池为空,则启动推测生成;否则,继续向量搜索直到请求池为空。证明过程考虑了四种情况:(1) 池为空且 $D_i$ 是最终结果:立即启动推测生成可利用空闲LLM引擎降低延迟。(2) 池非空且 $D_i$ 不是最终结果:推迟推测更有效,因为插入会延迟池中请求且计算是不必要的。(3) 池非空且 $D_i$ 是最终结果:立即启动允许LLM调度器进行更有效的优化。(4) 池为空且 $D_i$ 不是最终结果:不正确的推测一旦发现就会被终止,仅利用空闲GPU资源,推迟或启动效率相同。

通用系统的动态推测流水线算法:通用的RAG系统更为复杂,具有更大的LLM Batch Size和并行的向量搜索。此外,当请求与其他请求批处理时,我们无法立即终止推测生成。基于定理5.1,我们在算法2中设计了动态推测流水线策略。其主要思想是:仅当检索到的文档发生变化,且待处理的LLM请求数量低于预先确定的预填充迭代最大Batch Size(max_prefill_bs)时,才启动推测生成。该策略在当前LLM迭代后终止不正确的推测生成,这不影响批次中的其他请求。

function DYNAMIC_SPECULATIVE_PIPELINING(request):
    D = []
    while the vector search of request is not finished:
        # Produce the candidate documents at the next stage
        D_temp = VECTOR_SEARCH(request, D)
        if D_temp != D:
            if {request, D} in pool:
                Terminate {request, D} after the current iteration
            if pool.size < max_prefill_bs:
                pool.insert({request, D_temp})
            D = D_temp

系统实现细节:我们使用约5000行C++和Python代码实现了RAGCache系统原型。实现基于先进的LLM服务系统vLLM v0.3.0【索引编号1】。我们扩展了其在Pytorch【索引编号5,Pytorch: An imperative style, high-performance deep learning library+2019+NeurIPS】和Triton【索引编号6,Triton: an intermediate language and compiler for tiled neural network computations+2019+MAPL】中的预填充内核,以支持不同注意力机制(如多头注意力【索引编号7,Attention is all you need+2017+NeurIPS】和分组查询注意力【索引编号8,GQA: Training generalized multi-query transformer models from multi-head checkpoints+2023+arXiv】)的前缀缓存。

流水线向量搜索实现:我们在广泛使用的开源向量数据库Faiss【索引编号9,Pinecone: Introduction to Facebook AI Similarity Search (Faiss)+2024+URL】之上实现了动态推测流水线,并使其适配两种索引:IVF【索引编号10,The inverted multi-index+2014+IEEE TPAMI】和HNSW【索引编号11,Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs+2018+IEEE TPAMI】。对于IVF,我们将搜索拆分为多个阶段,每个阶段搜索部分聚类并返回当前的Top-$k$向量。对于HNSW,我们测量特定配置的平均搜索时间,将其拆分为更小的时间片,在每个时间片后返回当前的Top-$k$向量。

容错机制:RAGCache实现了两种容错机制来处理GPU故障和请求处理故障。GPU内存作为第一级缓存,存储知识树上层节点的KV Cache。由于LLM推理的前缀敏感性,GPU故障将使下层节点失效,从而导致整棵树失效。我们将部分最频繁访问的上层节点(如系统提示词)复制到主机内存中以实现快速恢复。我们还采用了超时机制来重试失败的请求。如果请求在完成第一次迭代前失败,将被重新计算;否则,请求可以通过重用存储的KV Cache继续计算。

实验环境

硬件配置:主要实验在AWS EC2 g5.16xlarge实例上进行。每个实例配置64个vCPU(AMD EPYC 7R32)、256 GiB主机内存、25 Gbps网卡,以及1张配备24 GiB显存的NVIDIA A10G GPU(通过PCIe 4.0 x16连接)。对于大型模型实验,使用两张NVIDIA H800 80GB GPU,通过NVLink互连,并通过PCIe 5.0 x16连接到主机。缓存的主机内存大小默认设置为192 GiB(单卡环境)或384 GiB(双卡环境)。

软件与模型配置
- 模型:Mistral-7B、LLaMA2-7B、Mixtral-8x7B(MoE模型)、LLaMA2-70B。大型模型采用张量并行和专家并行部署。
- 检索与数据集:使用基于Wikipedia语料库生成的文档数据集【索引编号12,Wikipedia (en) embedded with http://cohere.ai multilingual22-12 encoder+2024+URL】作为知识库。向量搜索使用IVF索引,1024个聚类,默认Top-$k$为2。评测负载使用两个代表性QA数据集:MMLU【索引编号13,Measuring massive multitask language understanding+2020+arXiv】和Natural Questions【索引编号14,Natural questions: a benchmark for question answering research+2019+TACL】。
- 基线系统:对比了vLLM【索引编号1】和SGLang【索引编号15,Efficiently programming large language models using sglang+2023+arXiv】。所有基线配置与RAGCache相同的模型并行度、最大Batch Size和向量数据库设置。

实验结果

1. 整体性能表现
- 实验内容:在MMLU和Natural Questions数据集上,使用Mistral-7B和LLaMA2-7B模型,评估不同请求率下的平均TTFT(Time-to-First-Token)和吞吐量。
- 实验结果:如图13和图14所示,在相同请求率下,RAGCache相较于vLLM将平均TTFT降低了 $1.2-4\times$,相较于SGLang降低了 $1.1-3.5\times$。在吞吐量方面,RAGCache比vLLM高出 $1.3-2.1\times$,比SGLang高出 $1.2-1.8\times$。
- 分析结论:性能提升归功于RAGCache利用GPU和主机内存缓存热点文档的KV Cache,避免了频繁的重计算。Mistral-7B上的性能差距大于LLaMA2-7B,因为LLaMA2-7B的KV Cache大小是前者的 $4\times$,导致命中率相对较低。MMLU比Natural Questions从文档缓存中获益更多。
图13 MMLU上的整体性能
图14 Natural Questions上的整体性能

2. 案例研究:不同的Top-k值
- 实验内容:在MMLU和Mistral-7B下,评估Top-$k$为1、3、5时的性能。
- 实验结果:如图15所示,RAGCache在各种Top-$k$值下,平均TTFT优于vLLM $1.7-3.1\times$,优于SGLang $1.2-2.5\times$。
- 分析结论:尽管随着Top-$k$增加,文档排列呈阶乘级增长,但由于知识树始终驱逐距离根节点最远的节点,确保了最常用的前缀保留在缓存中,RAGCache依然保持了优势。
图15 不同Top-k值下的性能

3. 案例研究:大型模型
- 实验内容:在两张H800 GPU上部署Mixtral-8x7B和LLaMA2-70B,评估不同请求率下的性能。
- 实验结果:如图16所示,在低请求率下,RAGCache将平均TTFT降低了 $1.4-2.1\times$。在高请求率下,vLLM无法满足SLO,而RAGCache仍能保持TTFT低于1.4秒。RAGCache的TTFT也比SGLang低 $1.2-2.6\times$。
- 分析结论:即使在大模型和更充足的GPU内存环境下,RAGCache的多级缓存策略依然展现出卓越的可扩展性和性能优势。
图16 大型模型下的性能

4. 消融实验:PGDSF策略
- 实验内容:对比PGDSF、原生GDSF、LRU和LFU替换策略在不同主机内存大小(8GiB到128GiB)下的命中率和TTFT。
- 实验结果:如图17和表2所示,PGDSF实现了最高的命中率,相较于GDSF提升了 $1.02-1.32\times$,相较于LRU提升了 $1.06-1.62\times$,相较于LFU提升了 $1.06-1.75\times$。平均TTFT也比基线策略低 $1.05-1.29\times$。
- 分析结论:PGDSF成功捕获了不同文档前缀的大小变化、访问模式和重计算成本差异,从而显著提升了缓存命中率。
图17 缓存替换策略消融实验

5. 消融实验:缓存感知重排序
- 实验内容:在饱和请求队列下(MMLU 2.5 req/s,NQ 1.4 req/s),评估缓存感知重排序的影响。
- 实验结果:如图18所示,RAGCache通过缓存感知重排序将平均TTFT降低了 $1.2-2.1\times$。
- 分析结论:在高请求率下,重排序策略有效减少了缓存颠簸,提升了系统效率。
图18 缓存感知重排序消融实验

6. 消融实验:动态推测流水线
- 实验内容:对比有无动态推测流水线在不同向量搜索比例($12.5\%$ 到 $100\%$)下的性能。
- 实验结果:如图19和表3所示,动态推测流水线使RAGCache的TTFT降低了高达 $1.6\times$,并将非重叠向量搜索时间减少了 $1.5-4.3\times$。
- 分析结论:推测流水线成功掩盖了大部分检索延迟,证明了其在重叠检索与生成步骤上的有效性。
图19 推测流水线消融实验

7. 调度时间开销
- 实验内容:测量RAGCache的调度时间(知识树查找更新、重排序、推测决策)。
- 实验结果:如表4所示,在各种请求率下,调度时间均保持在一毫秒以下。
- 分析结论:与秒级的TTFT相比,RAGCache的调度开销可以忽略不计。

补充细节

每个输出Token时间(TPOT):除了首个Token时间(TTFT),TPOT对LLM服务也至关重要。RAG由于增加了检索文档,显著延长了输入长度,从而导致预填充阶段的延迟(即TTFT)成为主要问题。RAGCache通过缓存最频繁检索文档的KV Cache降低了TTFT,并且由于预填充迭代的加速,它同样能够降低TPOT。

大Top-k值的影响:随着Top-$k$值的增加,文档排列的数量呈阶乘级爆炸式增长,使得它们被重用的可能性降低。RAGCache通过缓存具有较小Top-$k$值的文档(例如,为请求Top-5文档的查询缓存Top-3文档的KV Cache)来缓解这一问题,从而在命中率和缓存效率之间取得平衡。

相关工作(KV Cache重用):近期的研究【索引编号15;索引编号16,Prompt cache: Modular attention reuse for low-latency inference+2023+arXiv;索引编号17,CacheGen: Fast Context Loading for Language Model Applications+2023+arXiv;索引编号18,ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition+2024+arXiv】提出了在请求间重用KV Cache以减少冗余计算。然而,Prompt Cache和CacheGen可能会生成不准确的响应;SGLang和ChunkAttention主要识别GPU内存中可重用的KV Cache。RAGCache则利用了RAG的特定检索模式,构建了多级缓存系统,在保持生成结果完全准确的同时实现了更高的性能。

结论

本文提出了RAGCache,这是一个专门为检索增强生成(RAG)定制的多级缓存系统。基于对RAG系统的详细特征分析,RAGCache采用了一棵带有前缀感知替换策略的知识树来最小化冗余计算,并引入了动态推测流水线机制以在RAG工作流中重叠知识检索和LLM推理步骤。在各种模型和工作负载上的评估表明,RAGCache的性能显著优于当前最先进的解决方案(集成Faiss的vLLM),其TTFT降低了高达 $4\times$,吞吐量提升了高达 $2.1\times$。

方法细节引用汇总

以下是方法与系统实现细节中引用的参考文献汇总:
1. 【索引编号1,Efficient memory management for large language model serving with pagedattention+2023+SOSP】:用于说明RAGCache遵循vLLM的设计,将键值张量存储在非连续的内存块中以实现KV Cache重用,并作为系统原型的基础框架。
2. 【索引编号2,Improving WWW proxies performance with greedy-dual-size-frequency caching policy+1998】:用于说明PGDSF缓存替换策略是基于该经典的GDSF策略演变而来。
3. 【索引编号3,Improving approximate nearest neighbor search through learned adaptive early termination+2020+SIGMOD】:用于支撑动态推测流水线的设计洞察,即最终的Top-$k$文档可能在检索步骤的早期就已经出现。
4. 【索引编号4,Fundamentals of queueing theory+2018】:用于在简化分析中将LLM引擎建模为G/G/1队列,以推导推测流水线的最优策略。
5. 【索引编号5,Pytorch: An imperative style, high-performance deep learning library+2019+NeurIPS】:RAGCache扩展预填充内核所依赖的深度学习库。
6. 【索引编号6,Triton: an intermediate language and compiler for tiled neural network computations+2019+MAPL】:RAGCache扩展预填充内核所依赖的编译器语言。
7. 【索引编号7,Attention is all you need+2017+NeurIPS】:用于说明RAGCache支持标准的多头注意力机制。
8. 【索引编号8,GQA: Training generalized multi-query transformer models from multi-head checkpoints+2023+arXiv】:用于说明RAGCache支持分组查询注意力机制。
9. 【索引编号9,Pinecone: Introduction to Facebook AI Similarity Search (Faiss)+2024+URL】:RAGCache实现动态推测流水线所基于的开源向量数据库。
10. 【索引编号10,The inverted multi-index+2014+IEEE TPAMI】:Faiss支持的IVF索引结构,RAGCache对其进行了流水线改造。
11. 【索引编号11,Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs+2018+IEEE TPAMI】:Faiss支持的HNSW图索引结构,RAGCache对其进行了流水线改造。
12. 【索引编号12,Wikipedia (en) embedded with http://cohere.ai multilingual22-12 encoder+2024+URL】:实验环境使用的外部知识库数据集。
13. 【索引编号13,Measuring massive multitask language understanding+2020+arXiv】:实验环境使用的MMLU评测负载。
14. 【索引编号14,Natural questions: a benchmark for question answering research+2019+TACL】:实验环境使用的Natural Questions评测负载。
15. 【索引编号15,Efficiently programming large language models using sglang+2023+arXiv】:作为对比的基线系统SGLang。
16. 【索引编号16,Prompt cache: Modular attention reuse for low-latency inference+2023+arXiv】:相关工作中提及的KV Cache重用技术。
17. 【索引编号17,CacheGen: Fast Context Loading for Language Model Applications+2023+arXiv】:相关工作中提及的KV Cache压缩与重用技术。
18. 【索引编号18,ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition+2024+arXiv】:相关工作中提及的GPU内存KV Cache识别技术。