文章标题: CacheBlend: 通过缓存知识融合实现检索增强生成(RAG)的快速大语言模型服务
作者/机构: Jiayi Yao (University of Chicago/CUHK Shenzhen), Hanchen Li (University of Chicago), Yuhan Liu (University of Chicago), Siddhant Ray (University of Chicago), Yihua Cheng (University of Chicago), Qizheng Zhang (Stanford University), Kuntai Du (University of Chicago), Shan Lu (Microsoft Research / University of Chicago), Junchen Jiang (University of Chicago)

A1 主要贡献

本文旨在解决一个核心挑战:当一个大型语言模型(LLM)的输入包含多个文本块时,如何快速地组合它们各自预计算的键值缓存(KV Caches),以达到与昂贵的完全预填充(full prefill)相同的生成质量。这个问题在检索增强生成(RAG)等应用中尤为突出,因为这类应用的输入通常由多个检索到的文本块作为上下文。

核心问题: 现有的KV缓存重用方法存在局限性。
1. 前缀缓存(Prefix Caching): 仅能重用作为输入前缀的文本块的KV缓存。在RAG这类包含多个上下文块的场景中,只有第一个块能被重用,导致加速效果有限。
2. 完全KV重用(Full KV Reuse): 虽然可以拼接多个非前缀文本块的预计算KV缓存,但它忽略了文本块之间至关重要的“交叉注意力”(cross-attention),即一个文本块中的令牌(token)对其前面文本块中令牌的注意力。这种忽略会导致生成质量显著下降,甚至产生错误答案。

研究目标: 实现一种兼具“完全KV重用”的速度和“完全预填充”的生成质量的KV缓存融合机制。

核心创新点 (CacheBlend):
1. 选择性KV重计算(Selective KV Recompute): 提出一种新方案,该方案重用所有预计算的KV缓存(无论是否为前缀),并通过选择性地重新计算一小部分令牌(通常少于15%)的KV值来部分更新每个重用的KV缓存。这种方法能够有效地恢复被“完全KV重用”忽略的交叉注意力信息,从而保证生成质量。
2. 流水线化(Pipelining)计算与加载: 将选择性KV重计算的开销与从存储设备中检索KV缓存的延迟进行流水线化处理。通过这种方式,重计算带来的微小延迟可以被KV缓存的加载延迟所隐藏,使得CacheBlend能够将KV缓存存储在容量更大但速度较慢的设备(如SSD)上,而不会增加推理延迟。
3. 系统设计与实现: 在vLLM之上实现了CacheBlend系统,包含一个智能的加载控制器,用于根据存储设备速度和模型特性动态选择最优的重计算比例,以平衡延迟和生成质量。

通过在三个不同规模的开源LLM和四个基准数据集上的实验,CacheBlend相比完全KV重计算,将首个令牌生成时间(TTFT)降低了2.2-3.3倍,推理吞吐量提高了2.8-5倍,且不牺牲生成质量。

A3 背景知识与动机

第2节 背景知识

LLM推理流程。大多数LLM服务使用Transformer架构 [13, 16, 52]。收到输入令牌后,LLM首先进入预填充(prefill)阶段,将令牌转换为键(K)和值(V)向量,即KV缓存。预填充后,LLM进入解码阶段,迭代地生成下一个令牌,并将新令牌的K和V向量附加到KV缓存中。

预填充阶段详解。预填充阶段逐层计算KV缓存。每层的输入令牌嵌入首先被转换为查询(Q)、键(K)和值(V)向量,其中K和V向量构成了该层的KV缓存。然后,LLM将Q和K向量相乘以获得注意力矩阵,该矩阵表示每个令牌对其前面所有令牌的注意力。经过归一化和掩码处理的注意力矩阵与V向量进行点积运算,其结果向量将通过多个神经网络层,以获得下一层的令牌嵌入。

前缀缓存下的预填充。当一个前缀的KV缓存可用时,预填充阶段只需计算前向注意力矩阵(后缀令牌与前缀令牌之间的注意力),这直接影响生成的令牌。

预填充的性能瓶颈。预填充阶段可能非常缓慢,尤其是在处理长输入时。例如,在单个A40 GPU上,处理一个4000个令牌的输入(RAG中的典型上下文长度 [33]),Llama-34B或Llama-70B的预填充可能需要3到6秒。这种延迟导致用户在看到第一个生成词之前需要等待很长时间。最近的研究也表明,预填充可能成为吞吐量的瓶颈,消除预填充阶段可以将LLM推理系统的吞吐量提高一倍 [60]。

第3节 动机

3.1 重用KV缓存的机遇

文本重用现象。近期系统试图通过利用一个观察来减轻预填充开销:在许多LLM用例中,相同的文本会在不同的LLM输入中被重复使用。这使得重用这些文本的KV缓存成为可能。文本重用在需要提供上下文以确保高质量和一致性响应的场景中尤其普遍。

具体场景举例
* 在一个使用LLM管理内部记录的公司中,“IT部门谁在上一次全体会议上提议使用RAG来增强客户服务X?”和“IT部门谁毕业于Y大学?”这两个查询,虽然看似不同,但都需要IT部门员工列表作为必要的上下文。
* 在一个总结Arxiv论文的LLM应用中,“Arxiv上有什么热门的RAG技术?”和“最近有哪些数据集被用来对Arxiv上与RAG相关的论文进行基准测试?”这两个查询,都需要近期关于RAG的Arxiv论文作为上下文。

上下文预填充的开销。由于重用的上下文通常比用户查询包含更多信息,因此对输入中“上下文”部分的预填充占据了预填充开销的大部分 [22, 33]。因此,理想的做法是存储和重用这些文本的KV缓存,以避免在它们被再次使用时产生预填充开销。

3.2 为什么前缀缓存是不够的?

前缀缓存的原理与优势。一些近期系统通过重用KV缓存来减少预填充延迟。例如,在前缀缓存中,一个可重用文本块的KV缓存被预计算一次,如果该文本块出现在LLM输入的前缀位置,那么这个预计算的KV缓存就可以被重用来避免对前缀的预填充。前缀缓存的优点是,前缀的KV缓存不受后续文本的影响,因此生成结果与完全KV重计算(不使用KV缓存)完全相同。vLLM [36]、SGLang [59]和RAGCache [33]等系统都采用了这种方法。

前缀缓存的局限性。前缀缓存的缺点也很明显。为了回答一个查询,像RAG这样的应用通常会在LLM输入前附加多个文本块以提供必要的不同上下文。因此,除了第一个块,所有其他块的KV缓存都无法被重用,因为它们不是LL-M输入的-前缀。

多源上下文的需求。以上一节中的查询为例,要回答“IT部门谁在上一次全体会议上提议使用RAG来增强客户服务X?”,需要来自多个来源的上下文,包括IT部门员工信息、服务X的信息以及全体会议的会议记录。同样,在Arxiv总结应用中,回答示例查询需要LLM阅读多篇近期的RAG相关Arxiv论文作为上下文。这些上下文主题各异,不太可能出现在同一个文本块中,而是在回答特定查询时才被组合使用。

多文本块对质量的提升。为了从经验上证明在LLM输入中包含多个文本块的必要性,我们使用了两个流行的多跳问答数据集Musique和2WikiMQA。我们遵循RAG的常见做法,将上下文分割成128个令牌的块,并为每个查询检索top-k相关的块。实验结果表明,随着检索到的文本块数量增加,生成质量(以F1分数衡量)显著提高,尽管包含过多块会因“迷失在中间”问题 [40, 54]而损害质量。

图2. 随着检索到的文本块数量增加,生成质量得到改善。
图2. 随着检索到的文本块数量增加,生成质量得到改善。

结论。简而言之,前缀缓存只能节省第一个文本块的预填充时间,因此当LLM输入包含更多文本块时,即使这些块是可重用的,其节省的开销也微乎其微。

3.3 为什么完全KV重用是不够的?

完全KV重用的方法。完全KV重用旨在解决前缀缓存的问题,该方法由PromptCache [24]首创。它通过使用缓冲区来拼接独立预计算的、复现文本块的KV缓存,以保持每个文本块的位置准确性。例如,要拼接块C1和C2的KV缓存,PromptCache需要在一个假设的输入上预计算C2的KV缓存,该输入在C2前有一个长度不小于C1的虚拟前缀。这样,即使C2不是前缀,我们也能正确保留其KV缓存的位置信息,尽管每个块的KV缓存可能需要被多次预计算。

核心问题:忽略交叉注意力。然而,即使保留了位置信息,一个更根本的问题是,非前缀文本块(如C2)的KV缓存忽略了该块与前置文本(如C1)之间的交叉注意力(cross-attention)。这是因为在预计算KV缓存时,前置文本是未知的。

忽略交叉注意力的后果。忽略交叉注意力可能导致错误的响应。图3展示了一个例子,用户查询“梅西在世界杯上比C罗多进了多少个球?”前面附加了两个球员职业生涯统计的文本块。使用完全预填充或前缀缓存,结果清晰且正确。而使用完全KV重用时,LLM会开始 rambling(胡言乱语)并且无法给出正确答案。

图3. 一个LLM输入示例,查询前附加了两个文本块。(b) 完全KV重计算速度慢但答案正确。(c) 完全KV重用因忽略了块间交叉注意力(图4)而给出错误答案。
图3. 一个LLM输入示例,查询前附加了两个文本块。(b) 完全KV重计算速度慢但答案正确。(c) 完全KV重用因忽略了块间交叉注意力(图4)而给出错误答案。

交叉注意力的可视化分析。为了理解原因,我们观察了注意力矩阵。图4可视化了完全预填充和完全KV重用生成的注意力矩阵。由于完全KV重用是分开预计算每个块,两个块之间的交叉注意力在预计算时完全丢失了。在这个例子中,第一个块包含梅西的进球数,第二个块包含C罗的,查询需要比较二者。忽略这两个块之间的交互(交叉注意力)会导致错误的答案。

图4. (a) 完全KV重计算和 (b) 完全KV重用的注意力矩阵对比。黄色框突出显示了交叉注意力。右侧图表显示了由此产生的前向注意力矩阵的差异,这些差异是由于两种方法之间交叉注意力的不同造成的。
图4. (a) 完全KV重计算和 (b) 完全KV重用的注意力矩阵对比。黄色框突出显示了交叉注意力。右侧图表显示了由此产生的前向注意力矩阵的差异,这些差异是由于两种方法之间交叉注意力的不同造成的。

交叉注意力的普遍性。完全KV重用中交叉注意力的缺失导致了前向注意力矩阵(上下文令牌与最后几个令牌之间的注意力,直接影响生成结果)的显著差异。为了展示交叉注意力在多块LLM输入中的普遍性,图2对比了完全KV重计算(有交叉注意力)和完全KV重用(无交叉注意力)的响应质量。随着相关块数的增加,两者之间的质量差距变得更加明显,因为块数越多,输入不同部分之间的交叉引用和相互依赖性(即交叉注意力)就越多。

A2 方法细节

第4节 快速KV缓存融合

鉴于完全KV重计算速度太慢,而完全KV重用质量低下,我们追求的目标是:当一个LLM输入包含多个重用文本块时,如何快速更新预计算的KV缓存,使得生成的前向注意力矩阵(以及最终的输出文本)与完全KV重计算产生的差异最小。为此,我们提出了CacheBlend,它在每一层选择性地重计算一部分令牌的KV,同时重用其他令牌的KV。

4.1 术语

符号定义。表1总结了本节使用的符号。对于一个包含n个文本块的列表,我们用 $KV_{full}$ 表示完全KV重计算得到的KV缓存,用 $KV_{pre}$ 表示预计算的KV缓存,用 $KV_{new}$ 表示经CacheBlend更新后的KV缓存。KV缓存的第 $i$ 层 $KV_i$ 会产生前向注意力矩阵 $A_i$。

表1. 术语总结
表1. 术语总结

两种偏差的定义
* KV偏差(KV deviation): 对于一个KV缓存 $KV$ 在第 $i$ 层、令牌 $t$ 上的KV偏差,定义为其与 $KV_{full}$ 在该位置的绝对差值,记为 $\Delta_{kv}(KV_i, KV_{full_i})[t]$。它衡量了给定KV在特定令牌和层上与完全预填充KV缓存的差异程度。
* 注意力偏差(Attention deviation): 对于第 $i$ 层的前向注意力矩阵 $A_i$,其注意力偏差定义为它与 $A_{full_i}$ 差异的L2范数,记为 $\Delta_{attn}(A_i, A_{full_i})$。完全KV重用正是因为交叉注意力的缺失而导致了前向注意力矩阵的偏差。

目标重述。利用这些术语,我们的目标可以表述为:如何快速地将预计算的KV缓存 $KV_{pre}$ 更新为新的KV缓存 $KV_{new}$,从而使任意层 $i$ 上的注意力偏差 $\Delta_{attn}(A_{new_i}, A_{full_i})$ 最小化。

4.2 选择性地重计算KV缓存

工作流程。假设我们已经在每一层选定了一个要重计算的令牌子集(选择方法见§4.3),CacheBlend按以下步骤重计算这些选定令牌的KV(如图5(b)所示):
1. 首先,在每层 $i$ 的输入上应用一个掩码,将其缩减为选定的令牌子集。
2. 然后,将缩减后的输入转换为 $K_i, V_i$ 和 $Q_i$ 向量,这些向量也将仅限于选定的令牌。
3. 接着,通过重用第 $i$ 层上未选定令牌的KV缓存条目来扩展 $K_i$ 和 $V_i$ 向量,这样注意力矩阵就能包含选定令牌与所有其他令牌之间的注意力。
4. 最后,运行相同的注意力模块,生成下一层的输入。

图5. (a) 完全KV重计算与 (b) 选择性KV重计算在单层上的对比图示。
图5. (a) 完全KV重计算与 (b) 选择性KV重计算在单层上的对比图示。

计算开销。这些改动对Transformer的具体流程假设很少,可以集成到许多流行的Transformer模型中(详见§6)。重要的是,计算开销与选定令牌的数量成正比。如果我们每层重计算 $r\%$ 的令牌,总计算开销将是完全预填充的 $r\%$。

4.3 选择要重计算的哪些令牌

基本直觉。为了减少由KV偏差导致的注意力偏差,我们的直觉是优先重计算那些具有高KV偏差的令牌的KV。

高KV偏差令牌的效果验证。图6展示了在Musique数据集上,当我们选择并重计算具有最高KV偏差的 $r\%$ 的令牌 $t$ 时,所有层 $i$ 的平均注意力偏差 $\Delta_{attn}(A_i, A_{full_i})$ 的变化。随着重计算比例($r$)的增加,注意力偏差逐渐减小,并且最大的降幅发生于重计算具有最高KV偏差的头几个令牌时。这表明了一个关键洞见。

图6. 随着每层重计算更多令牌,注意力偏差减小。重要的是,注意力偏差的最大降幅来自于重计算具有最高KV偏差(即HKVD令牌)的令牌。
图6. 随着每层重计算更多令牌,注意力偏差减小。重要的是,注意力偏差的最大降幅来自于重计算具有最高KV偏差(即HKVD令牌)的令牌。

洞见1。在第 $i$ 层,重计算具有较高KV偏差(即 $\Delta_{kv}(KV_i, KV_{full_i})[t]$)的令牌 $t$ 的KV,能更大幅度地减少注意力偏差(即 $\Delta_{attn}(A_i, A_{full_i})$)。因此,我们应该选择每层中具有最高KV偏差的令牌进行重计算,我们称之为高KV偏差(High-KV-Deviation, HKVD)令牌。

问题1:需要重计算多少令牌?。在§7中,我们凭经验表明,选择10-20%的令牌作为HKVD令牌并重计算它们的KV,足以大幅减少注意力偏差并保持生成质量。这可以通过注意力稀疏性来解释,即在注意力矩阵中,高注意力通常只发生在少数令牌与其前序令牌之间 [14, 15, 43, 58]。图7显示了单层上KV偏差的分布,大约10-15%的令牌具有比其他令牌高得多的KV偏差,这证实了交叉注意力的稀疏性。

图7. 单层上不同令牌的KV偏差分布。
图7. 单层上不同令牌的KV偏差分布。

问题2:如何在没有真实KV值的情况下识别HKVD令牌?。朴素地识别HKVD令牌需要知道完全重计算的 $KV_{full_i}$,这违背了选择性重计算的目的。我们观察到不同层上的HKVD令牌并非独立:

洞见2。在一层上具有最高KV偏差的令牌,很可能在下一层也具有最高的KV偏差。
图8显示了相邻两层之间令牌KV偏差的Spearman秩相关分数,结果显示不同层之间HKVD令牌存在持续的高度相似性。这种相关性的直觉在于,Transformer模型中每个令牌的输入嵌入在层与层之间变化缓慢 [44, 47],因此由输入嵌入线性变换生成的KV缓存也应具有相似性。

图8. 相邻两层之间每个令牌KV偏差的秩相关性。
图8. 相邻两层之间每个令牌KV偏差的秩相关性。

渐进式过滤方案。基于HKVD令牌在层间的相关性,一个简单的方案是在第一层进行预填充,选出HKVD令牌,然后在所有后续层只更新这些令牌的KV。然而,仅使用第一层的注意力偏差来选择所有层的HKVD令牌可能不够可靠。因此,我们采用一种渐进式过滤方案(如图9所示):
1. 假设我们平均每层要选择 $r\%$ 的HKVD令牌。
2. 在第一层,我们根据令牌级的注意力偏差选择 $r_1\%$ 的令牌($r_1$ 略高于 $r$),并将它们作为第二层的HKVD令牌。
3. 在第二层,我们重计算这 $r_1\%$ 的HKVD令牌的KV,并从中选出具有最高令牌级注意力偏差的 $r_2\%$ 的令牌($r_2$ 略小于 $r_1$),作为下一层的HKVD令牌。
4. 依此类推。这种方案凭经验能更可靠地识别出在多个层上都具有高注意力偏差的HKVD令牌。
执行HKVD计算的第i层的KV缓存空间虽然同时持有更新后的KV和预计算的KV,但一旦推理进行到第i+1层,第i层额外的预计算KV会立即被丢弃,使得HKVD的内存开销可以忽略不计。

图9. CacheBlend通过仅计算前一层选定的HKVD令牌的KV偏差,并从中选择具有高KV偏差的令牌,来选定当前层的HKVD令牌。
图9. CacheBlend通过仅计算前一层选定的HKVD令牌的KV偏差,并从中选择具有高KV偏差的令牌,来选定当前层的HKVD令牌。

第5节 CacheBlend系统设计

基本洞见。如果选择性KV重计算(§4.3)的延迟比将KV加载到GPU内存的延迟要快,那么通过适当地将选择性KV重计算和KV加载进行流水线化,KV重计算的额外延迟可以忽略不计。

流水线化KV加载与重计算。在CacheBlend中,一层的选择性重计算可以在前一层的预计算KV缓存被加载到GPU后立即开始。因此,如果加载一层预计算KV的时间长于或等于选择性KV重计算的时间,KV加载的延迟就能隐藏选择性重计算的延迟。
例如,对于Llama-7B模型和4K长度的上下文,重计算15%的令牌每层仅需3毫秒,而从NVME SSD加载一层KV缓存需要16毫秒(§7)。在这种情况下,KV加载可以完全隐藏重计算的延迟。相反,对于Llama-70B,重计算15%的令牌需要7毫秒,而从NVME SSD加载一层KV仅需4毫秒,此时加载延迟不能完全隐藏重计算延迟。因此,需要一个控制器来智能地选择重计算比例以及KV缓存的存储位置。

图10. (a) 巧妙选择重计算比例不会产生额外延迟。(b) 巧妙选择存储设备来存储KV可以在不增加延迟的情况下节省成本。
图10. (a) 巧妙选择重计算比例不会产生额外延迟。(b) 巧妙选择存储设备来存储KV可以在不增加延迟的情况下节省成本。

5.1 关键组件

加载控制器(Loading Controller)。该控制器解决两个设计问题:
1. 给定存储设备,如何选择重计算比例? 控制器使用两个延迟估计器来找到一个理想的重计算比例 $r\%$,使得重计算延迟接近加载延迟。它计算重计算延迟 $T_{recompute}(r\%, L_{ctx}, M)$ 和加载延迟 $T_{load}(L_{ctx}, M, S_{dev})$,然后选择使 $T_{recompute}$ 等于 $T_{load}$ 的 $r\%$,并取该值与一个保证质量的最小重计算比例 $r^*\%$(实践中为15%)中的最大值。
2. 给定重计算比例,如何选择存储设备? 控制器可以使用一个存储成本估计器 $C_{store}(L_{ctx}, M, T_{dur}, S_{dev})$,并结合延迟估计器,在所有满足 $T_{recompute} \ge T_{load}$ 的设备中,选择成本最低的一个。

KV缓存存储(KV cache store)。该组件将LLM输入分割为多个可重用或新的文本块。每个块通过哈希找到其对应的KV缓存,实现方式类似于vLLM [36]中的块哈希。新的KV缓存被添加到存储设备中。当存储设备满时,采用最近最少使用(LRU)策略进行淘汰。

融合器(Fusor)。缓存融合器(§4)通过选择性重计算来合并预计算的KV缓存。它逐层工作,等待上一层的重计算完成并且当前层 $i$ 的KV缓存被加载到GPU内存中的队列后,才使用加载控制器计算出的重计算比例 $r\%$ 进行选择性重计算。

5.2 整合流程

系统工作流。图11展示了CacheBlend在一个LLM推理工作流中的整合方式。
1. 用户提交问题后,系统检索一系列相关的文本块。
2. 加载控制器查询KV缓存管理器,确认这些文本块的KV缓存是否存在及其存储位置。
3. 加载控制器计算出理想的选择性重计算比例,并将其发送给融合器,同时开始将KV缓存加载到GPU内存的队列中。
4. KV缓存融合器持续地对队列中的KV缓存进行逐层重计算,直到所有层处理完毕。
5. 最后,融合后的KV缓存被输入到LLM推理引擎,用于生成用户问题的答案。

图11. 在单个请求的LLM上下文增强生成流程中,CacheBlend系统(绿色星标部分)的作用。CacheBlend使用检索器提供的文本,与存储设备交互,并在LLM推理引擎之上提供KV缓存。
图11. 在单个请求的LLM上下文增强生成流程中,CacheBlend系统(绿色星标部分)的作用。CacheBlend使用检索器提供的文本,与存储设备交互,并在LLM推理引擎之上提供KV缓存。

A4 实验

实验环境

  • 模型架构:
    • Mistral-7B [30]
    • Yi-34B [56](使用8位量化)
    • Llama-70B [2](使用8位量化)
  • 硬件配置:
    • 平台: Runpod [10]
    • GPU: 2 x NVIDIA A40
    • CPU RAM: 128 GB
    • 存储: 1TB NVME SSD (实测吞吐量 4.8 GB/s)
    • GPU使用: Mistral-7B和Yi-34B使用1个GPU,Llama-70B使用2个GPU。
  • 软件配置:
    • 实现: 基于 vLLM [36]
    • 语言/库: Python, PyTorch v2.0
  • 数据集:
    • 2WikiMQA [27]: 多跳问答数据集,测试推理能力,包含200个测试用例。
    • Musique [51]: 多文档问答数据集,测试多跳推理能力,包含150个测试用例。
    • SAMSum [25]: 对话摘要数据集,测试少样本学习能力,包含200个测试用例。
    • MultiNews [20]: 新闻文章摘要数据集,包含60个抽样用例。
    • Musique extended / 2WikiMQA extended: 为模拟RAG场景中的块重用而创建的合成数据集。从原数据集中随机选取1500个查询,用GPT-4 API为每个查询生成3个相似查询,共6000个查询。上下文被切分为512令牌的块,并为每个查询检索top-6相关的块。

实验结果

  • 降低TTFT与质量保持:

    • 实验内容: 比较CacheBlend与完全KV重计算、前缀缓存、完全KV重用在多个模型和数据集上的TTFT和生成质量(F1分数或Rouge-L分数)。
    • 实验结果 (图12): 与完全KV重计算和前缀缓存相比,CacheBlend的质量下降在0.02以内,但TTFT显著降低了2.2-3.3倍。虽然比完全KV重用慢,但其质量稳定地大幅超越后者(在许多情况下超过2倍)。
    • 与RAG方法对比 (图13): 与MapReduce相比,CacheBlend的TTFT低2-5倍,F1分数更高。

    图12. CacheBlend相较于完全KV重计算,在四个数据集和三个模型上将TTFT降低了2.2-3.3倍,且质量下降可忽略不计。
    图12. CacheBlend相较于完全KV重计算,在四个数据集和三个模型上将TTFT降低了2.2-3.3倍,且质量下降可忽略不计。

    图13. CacheBlend与MapReduce和MapRerank在Yi-34B上的生成质量对比。
    图13. CacheBlend与MapReduce和MapRerank在Yi-34B上的生成质量对比。

  • 提高吞吐量:

    • 实验内容: 在合成的RAG数据集上,比较不同请求率下各方法的延迟和吞吐量。
    • 实验结果 (图14): CacheBlend在所有模型和数据集上,实现了比所有基线方法低2.8-5倍的延迟和更高的吞吐量。它比完全KV重计算快是因为只重计算少量令牌;比完全KV重用质量高是因为恢复了交叉注意力;比前缀缓存性能好是因为前缀缓存对同一数据块需要存储多个不同前缀的版本,导致存储效率低和缓存未命中率高。

    图14. 在RAG场景中,与质量相似的基线相比,CacheBlend以更低的TTFT实现了更高的吞吐量。
    图14. 在RAG场景中,与质量相似的基线相比,CacheBlend以更低的TTFT实现了更高的吞吐量。

  • 敏感性分析:

    • 不同块数和块长 (图15a, 15b): 在不同块数和块长设置下,CacheBlend的计算时间节省比例保持相似。
    • 不同重计算比例 (图16): 在所有数据集上,使用5%-18%的重计算比例,CacheBlend的质量损失(F1或Rouge-L分数)相较于完全KV重计算最多只有0.002。这个比例能带来4.1-6.6倍的TTFT降低。
    • 不同批处理大小 (图15c): 随着批处理大小增加,预填充开销在总延迟中占比更大,因此CacheBlend对预填充阶段的改进对整体延迟降低的效果更显著。
    • 不同存储设备 (图17): 无论KV缓存存储在RAM还是较慢的SSD上,CacheBlend都能持续降低TTFT并保持质量。对于较慢的存储,CacheBlend与完全KV重用的延迟差距变小,因为总延迟更多地由加载延迟主导。

    图15. CacheBlend在不同块数、块长和批处理大小下均优于基线。
    图15. CacheBlend在不同块数、块长和批处理大小下均优于基线。

    图16. 在Yi-34B上,使用5%–18%的选择性重计算比例,CacheBlend与完全KV重计算相比质量损失极小。
    图16. 在Yi-34B上,使用5%–18%的选择性重计算比例,CacheBlend与完全KV重计算相比质量损失极小。

    图17. 当使用RAM和较慢的磁盘时,CacheBlend的表现优于基线。
    图17. 当使用RAM和较慢的磁盘时,CacheBlend的表现优于基线。

A5 结论

本文提出了CacheBlend,一个用于组合多个预计算KV缓存的系统,这些KV缓存对应的文本在LLM输入中被拼接在一起。为了保持生成质量,CacheBlend通过选择性地重计算一小部分令牌的KV缓存值来恢复这些文本间的交叉注意力。通过在四个数据集和三个模型上的实验证明,与完全KV重计算相比,CacheBlend在质量下降可忽略的情况下,将首个令牌生成时间(TTFT)降低了2.2-3.3倍,并将吞吐量提高了2.8-5倍。

A6 附录

A N维位置恢复

N维空间中的RoPE定义。我们首先给出N维空间中旋转位置编码(RoPE)的定义。
定义1 (旋转位置编码, ROPE[50])。令向量 $\mathbf{q}, \mathbf{k} \in \mathbb{R}^d$ 分别表示需要在某个位置 $m$ 进行嵌入的查询向量和键向量,记为 $\mathbf{q}_m, \mathbf{k}_m \in \mathbb{R}^d$。RoPE将位置信息编码如下:

公式1
公式1

其中旋转矩阵为:
公式2
公式2

超参数 $\Theta = \{\theta_i = 10000^{-2i/d}, i \in [0, 1, ..., \frac{d}{2}-1]\}$。

位置恢复方法有效性的证明。我们的位置恢复方法之所以有效,是因为一对令牌之间的注意力分数与其绝对位置无关。以下是该不变性的证明。
命题A.1 (RoPE仅依赖于相对位置)。令向量 $\mathbf{k} \in \mathbb{R}^d$ 表示一个键向量,$\mathbf{k}_m \in \mathbb{R}^d$ 表示在固定位置m嵌入的键向量。令向量 $\mathbf{q} \in \mathbb{R}^d$ 表示一个查询向量,$\mathbf{q}_{m+l} \in \mathbb{R}^d$ 表示在位置(m+l)嵌入的查询向量。那么注意力分数 $\mathbf{q}_{m+l}^T \mathbf{k}_m$ 推导如下:

公式3
公式3

其中 $\{\mathbf{q}, \mathbf{k}\}[i]$ 表示向量 $\{\mathbf{q}, \mathbf{k}\}$ 的第i个元素,$h_i$ 表示点积 $\mathbf{q}[i]^T\mathbf{k}[i]$。注意力分数 $\mathbf{q}_{m+l}^T \mathbf{k}_m$ 仅依赖于相对距离 $l$,而不是绝对位置 $m$。