前言
在本专栏之前的文章中，我们已经了解过一些Cache相关的内容：比如KVCache、PageAttention、GPTCache等。这些Cache基于推理特点或业务特点对关键数据进行Cache, 从而加速端到端应用的生成过程。本文将介绍一个最新的Cache：RAGCache，它是针对RAG应用的特点设计了一套分层的缓存系统，通过它可以提升应用的吞吐量，加速了应用的处理。
背景
RAG 基本流程
为了最小成本的减少模型幻觉，给模型注入新的知识，业界目前最通用的方法是使用RAG技术来对模型训练时未学习的知识进行注入。大致流程如下（更详细的介绍参考之前的文章）：
LLM 推理过程
LLM的推导一般分为两个步骤：prefill和decode。在prefill阶段，会对输入的token计算KV tensor（如图所示, K matrix和V matrix中的每个行向量表示一个token的语义表示）。之后，在decode阶段，会根据之前的KV向量信息不断推导next token的。可以看到，在prefill阶段是非常耗时的，因为他需要为请求中的所有token生成向量。请求越长，处理时间越长。尤其在RAG的场景下，会携带外部检索到的更多的文档信息，进一步增大了处理时间。
由于新知识注入，会将原问题进行扩充，从而导致LLM生成时引入长序列生成问题（序列越长，生成的处理时间越长），导致生成答案时的延迟增加。
优化思路
生成步骤的主要耗时主要在attention模块中处理长序列的KV tensor中，那一个自然而然的优化想法就是：把之前检索过的文档的KV tensor的信息缓存下来。 如果cache命中了，那么直接使用cache中的kv tensor信息，从而避免重新计算长序列的kv tensor信息。RAGCache系统就是基于此思想来工作的。
RAGCache
RAGCache是一种专为检索增强生成（RAG）量身定制的新型多级动态缓存系统。它的核心设计思想是知识树+PGDSF（Prefix-aware Greedy Dual-Size Frequency） 。考虑RAG应用特性，它将检索到的知识的中间状态（KV tensor) 组织在一棵树中，并利用PGDSF来将重要的KV tensor保留在缓存中。RAGCache的整体架构如下：
在RAGCache的作用下，一次完整的请求处理流程如下：
- 系统收到用户请求。
- RAGController从外部知识库中检索到相关的文档。
- RAGController将检索到的文档传到Cache Retriever去检索是否有对应缓存的KV tensor。
- 如果找到对应的KV tensor, 那么将对应kv tensor和请求交给LLM推理引擎去进行后续的生成。
- 如果没有找到对应的KV tensor(cache miss), 那么直接将文档内容和请求交给LLM推理引擎去生成回复。
- 在LLM生成第一个token之后，将已经生成的KV tensor交给RAGContoller去更新缓存。
接下来来看一下这个缓存系统中的一些主要模块。
Cache检索器（Cache Retriver)
检索器的主要职责是负责检索cache的KV tensor。实现的方式是将kv tensor组织成一棵树，每个节点上包含对应文档的KV tensor，从根节点到叶节点的路径就是一次请求的关联文档的所有KV tensor（文档的顺序信息在树的路径上得以保留，对后续LLM生成至关重要。由于GPU内存大小的限制，不能将所有的内容都存储在节点中，只能够存储一部分关键的kv tensor存储到GPU内存中用于计算，其他的swap out到主机内容中。如果确定不再使用，则从主机内存中free掉。
RAG 控制器（RAGController）
从上面的请求处理流程中可以看到，RAGController负责对请求、缓存的处理进行路由。RAG 控制器针对RAG流程进行了一些系统优化，主要包括PGDSF、重排序和pipline。其中PGDSF主要是用来尽可能的将重要的信息保存在缓存中，从而减小cache miss的概率。重排序主要是重新安排请求以提高缓存命中率并防止抖动，同时还能保证请求的公平性，减轻饥饿问题。pipline是一种动态流水线，用来使RAG的知识检索部分和大语言模型推理过程重叠，以减小延迟。
PGDSF
Prefix-aware Greedy-Dual-Size-Frequency (PGDSF, 翻译成中文是前缀感知贪婪双尺寸频率算法，感觉更难懂)是一种替换策略。RAG Cache需要一种策略来决定树节点在缓存中的位置，是位于更快访问到的GPU内存？还是位于访问速度慢一些但内存更大一些的GPU内存？又或者直接释放掉。本质上看上去有点类似于我们平时见到的缓存替换（比如redis缓存），那么自然而然的想法是：可以使用LRU等常见的替换策略。RAGCache中没有使用这些传统的策略，因为他们的评估标准有些单一（比如LRU按访问频率），可能不能将最有用的节点筛选出来保存在GPU内存中。相比而言，PGDSF基于访问频率、大小和访问花费等多个维度来进行评估，从而能够筛选中最有价值的节点。该方法的计算方式如下：
低优先级的节点被有限驱逐。这里的Clock跟踪节点访问的新近程度，在RAGController中为GPU和CPU分别维护一个逻辑时钟，从而适应cache层级。Frequency表示时间窗口内一个文档被检索的总次数。Cost是一个文档计算KV tensor所花费的时间。Size表示一个文档token化之后的token数量。
重排序（Cache-aware Reordering ）
重排序的背后的思想是：使用相同文档的不同请求可能不会同时到达，如果对不同的请求进行重排序，那么是由可能提高cache命中率，从而减小延迟。一个示例如下：比如q1,q3引用文档1， q2,q4引用文档2，那么处理他们的时候cache的，命中情况如下：
如果对他们进行重排序，那么cache的命中率如下，显著降低：
因此，引入重排序理论上是能够提升性能的。随之而来的问题是：重排序的依据是什么？论文中给出的计算依据是：
这里的Cached Length指针对这个问题需要cache的长度，Computation Length是指针对这个问题需要重新计算的长度（我理解是问题大小）。这样计算的依据如下（论文中描述的有些抽象，这里基于自己理解做了扩充）：
- 计算长度相同，cached Len越大越优先花费越小。
- cached长度相同，computation len越小越先计算花费越少
pipeline
pipeline的思路是：由于检索生成和模型推理生成都是比较耗时的过程，如果串行进行，在检索的时候，GPU其实是处于空闲状态的，没有利用上。为了能提升GPU的利用率，在检索过程中，第一次重排序之后，会立刻用当前排序的中间结果交给模型进行推理；在后续的重排序过程中，如果排序结果和之前不一致，则中断之前的推理的过程，开始使用新的排序结果来进行推理。这样检索和模型推理的过程几乎同时进行。如果每次排序结果都是一致的，那么当最终检索完成排序后，模型也同时得到推理的结果。通过这种方式来减少端到端的处理延迟。
在RAG Cache中，对以上pipeline进行了优化。在高负载的情况下，上面的一些错误计算会导致系统整体的性能降低。所以，优化的思路是：基于负载动态的开启pipeline。这里建模LLM engine包含一个G/G/1队列。执行的策略是：如果在一个阶段结束后(一次排序），文档序列和上一次的文档序列不同且请求池为空，则启动LLM推测；如果不是，则继续向量搜索，直到池为空。这里的证明参考原文档，这里不做展开了。
总结
本文从一些基本概念出发，逐步介绍了RAGCache的组成和各个模块中的一些关键设计和实现思路。相比于之前介绍过的其他Cache，RAGCache是专门针对RAG应用的特点而设计的缓存系统，其中一些优化思想还是非常值得去学习的。最终的优化效果数据以及一些实验结论在文中没有描述，感兴趣的同学可以自己去论文中获取。本文更侧重于原理的学习和理解。
参考文档
RAGCache: Efficient Knowledge Caching for Retrieval-Augmented Generation
