作者: Yinsicheng Jiang, Yeqi Huang, Liang Cheng, Cheng Deng, Xuan Sun, Luo Mai
核心问题与研究目标
当前,许多AI应用如检索增强生成(RAG)、AI记忆层和多智能体编排等越来越依赖于长上下文推理,导致预填充(prefill)延迟成为主要性能瓶瓶颈。现有的预填充加速技术面临一个权衡:要么为了保持推理质量而牺牲键值缓存(KV-cache)的重用率,要么为了提高重用率而牺牲推理质量。本文的研究目标是提出一种新的方法,在不牺牲模型准确性的前提下,显著降低长上下文输入的预填充延迟。
核心观察与创新点
本文观察到,在真实世界的长上下文工作负载中,无论是在同一对话的多个轮次之间,还是在领域特定应用的不同并行会话(用户查询)之间,上下文块(Context Blocks, CBs)都存在显著的重叠。基于这一观察,本文提出了CONTEXTPILOT,一个通过引入上下文重用机制来加速预填充的系统。
CONTEXTPILOT的核心贡献如下:
1. 上下文索引(Context Indexing):设计了一种索引机制,能够高效地跟踪跨并行会话和多轮对话历史中已缓存的上下文块。该索引支持通过以下方式快速检索先前存储的上下文:(i) 对齐传入上下文和已缓存块之间的前缀重叠部分;(ii) 遍历多轮历史中引用的块以恢复超出直接前缀匹配的可重用片段。该机制的构建和维护开销较低。
上下文排序(Context Ordering):提出了一种上下文排序算法,该算法查询索引以在会话内(跨轮次)和会话间选择和排列上下文块,明确目标是在固定的上下文预算下最大化缓存命中率。为了减轻重排序可能引入的准确性损失,引入了简洁的顺序标注(order annotations),这些标注保留了原始的相关性排名和结构线索,使LLM能够以与预期语义一致的方式解释重排序后的提示。
上下文去重(Context De-Duplication):通过查询索引,CONTEXTPILOT识别与已缓存上下文重叠的上下文块(包括部分重叠),并用指向其在提示(或先前轮次)中原始出现位置的简洁位置标注替换重复的跨度,从而在保持准确性的同时避免了冗余的预填充计算。
主要成果
大量的评估表明,CONTEXTPILOT相比现有技术(如CacheBlend, LMCache, RadixCache, RAGCache)能将LLM预填充延迟降低高达3倍,同时保持了推理质量。在更长的上下文长度下,由于其新颖的上下文标注设计,它甚至可以提高推理质量和答案准确性。特别是在大型MoE模型(如DeepSeek-R1)上,CONTEXTPILOT能将缓存命中率从5-6%提升至38-60%,在16和32个GPU上将预填充吞吐量提高了1.52-1.81倍。
长上下文推理系统概述。长上下文推理系统通过外部上下文块(如检索到的文档、块或记忆)来增强LLM的事实基础和推理能力。本文中,“上下文块”指任何注入到模型中的离散外部上下文单元,无论是检索的文档、文档块还是记忆条目。这一趋势由两种主流范式驱动:(1) 检索增强生成(RAG),它为每个查询从外部语料库中检索Top-K个最相关的文档,服务于在线延迟敏感服务和离线吞吐量导向的流水线。(2) AI记忆层系统(如Mem0),它动态地提取、整合和检索跨会话的用户特定记忆,并将相关的上下文块注入每个查询以实现个性化交互。在这两种情况下,推理引擎都执行预填充来编码这些上下文块,并执行解码来生成响应。
系统工作流程与前缀缓存。一个典型的系统(如图1所示)在会话和对话轮次之间交替进行检索(或记忆查找)和生成。在每一轮中,M个并发提示各自接收K个相关的上下文块。推理引擎对上下文进行编码并生成响应,这些响应连同更新后的对话历史一起进入下一步,从而实现高效的多轮推理。这些长上下文推理系统通常使用前缀缓存来提高预填充效率。基于Trie的实现(如【Zheng et al., 2024, Sglang: Efficient execution of structured language model programs】)将令牌分层组织,每个节点存储一个令牌序列及其KV缓存,通过单次遍历实现最长前缀匹配。另一种哈希表设计(如【Kwon et al., 2023, Efficient memory management for large language model serving with pagedattention】)则直接将完整的前缀映射到KV块标识符。
预填充延迟成为瓶颈。随着现代LLM要求不断扩大的上下文窗口,长上下文推理系统面临着一个关键的预填充延迟瓶颈。这由两个原因驱动:(1) 增加检索到的上下文块数量以扩大信息覆盖面,以及 (2) 通过检索完整文档或完整记忆历史并应用上下文工程方法来丰富上下文信息。
长上下文的收益与代价。对我们工作负载数据的分析显示,这两种方法都带来了显著的准确性提升。将检索参数k从较低值扩展到较高值可将准确性提高多达20%,而检索完整文档也取得了类似的性能改进。然而,扩展的上下文窗口(即更长的上下文块输入)会引入大量的预填充开销,甚至在超过一定长度后会降低推理质量。我们的跟踪数据显示,LLM推理引擎通常处理20k–130k的预填充令牌,在单个H100 GPU上执行32B密集模型时导致3–10秒的延迟。对于像混合专家(MoE)这样更大的模型,预填充延迟可能更高。因此,预填充成为主要的瓶颈,降低了用户体验,并阻碍了长上下文应用的广泛部署。
现有方法的局限性。为了解决日益增长的长上下文成本,现有的KV缓存重用方法存在一些问题:
精确前缀匹配导致低KV缓存重用率。现有的前缀缓存机制,如RadixCache、LMCache和RAGCache,严重依赖于精确的令牌级或文档级匹配。即使是微小的变动,如空格差异或轻微重排的令牌和文档,也会阻止重用。我们的评估(第7.1节)显示,尽管相关查询中检索到的文档有大量重叠,但缓存命中率仍然持续偏低。例如,在MultihopRAG数据集上使用Qwen3-32B模型,KV缓存命中率仅为4.6%。在NarrativeQA上使用Llama3.3-70B模型,命中率也只有5.5%,导致大部分缓存未被使用。
近似KV缓存匹配降低质量。为了改善低缓存命中率,最近的技术如CacheBlend采用了近似KV缓存匹配。它们不进行精确前缀匹配,而是测量KV值(浮点向量)的相似性,并在相似度超过经验设定的阈值时重用缓存状态。然而,KV值的相似性并不能可靠地指示缓存状态是否可以在不同上下文和请求之间重用。近似匹配会降低准确性,且误差会在多轮交互中累积。我们的评估(第7.1节)表明,在多个模型(如Qwen3-32B, Qwen3-4B, Llama3.3-70B)和数据集(如MultihopRAG, NarrativeQA, QASPER)上,近似匹配可能导致准确性下降9-11%(从约60%降至约50%),这使其无法部署在许多要求高保真度的服务中。
设计的核心动机。我们的设计基于一个关键观察:真实世界的长上下文工作负载在会话和对话轮次中都表现出上下文块的大量重叠:
跨会话重叠。(1) 图2a展示了多个用户查询同一个人不同方面时出现的重叠检索。尽管检索到的文档以不同的顺序出现以反映每个查询的相关性,但它们的内容在很大程度上是重合的。对MultihopRAG、NarrativeQA和QASPER数据集的跟踪研究证实了这一趋势:如图3所示,分别有79.2%、57.4%和49.6%的问题利用了访问最频繁的前20%文档,表明跨会话存在广泛的上下文共享。
多轮对话内的重叠。(2) 图2b显示,在多轮交互中,用户经常重新访问相关主题,导致检索器返回相同的文档,只是排名略有不同。由于之前的轮次成为输入的一部分,后续的检索经常会重复已经存在于缓存历史中的内容。我们的MT-RAG跟踪研究量化了这一效应:在任何一轮中,平均有40%的检索文档与同一会话中较早轮次的文档重叠。
重叠带来的重用机会。上下文块之间的显著重叠揭示了上下文重用的明确机会,从而提高KV缓存命中率。具体来说,我们确定了在真实世界长上下文应用中普遍出现的三种机会:
跨会话的上下文块排序提升KV缓存重用。(1) 如图2a所示,如果将第二和第三个用户的上下文块重新排序以匹配第一个用户的序列,所有三个上下文将共享一个相同的前缀,实现100%的KV缓存重用。基于跟踪数据的重排序实验证实了这一潜力。将上下文块顺序与前缀缓存结构对齐,在MultihopRAG、NarrativeQA和QASPER上分别将KV缓存命中率提高到38.9%、20.2%和16.5%,比基线利用率高3-8倍(第7.3节)。因此,策略性的上下文块重排序可以显著减少跨用户的冗余预填充计算。至关重要的是,这种重排序只带来微小的准确性损失:在相同数据集上仅为0.1-3.3%(第7.3节)。如表1所示,我们对DEmO排序研究的复现表明,现代LLM对输入顺序的敏感度远低于早期模型。微小的性能下降是因为为前缀优化的排序有时会将重要的上下文块移到列表的中间,使其受到“中间迷失”效应的影响。我们稍后将讨论恢复这种微小损失的策略。
多轮重叠的去重降低预填充成本。(2) 图2b显示,多轮检索常在对话轮次间返回重叠的上下文块。通过对这些块进行去重,只处理新内容和对话历史,可以大大减少预填充期间的上下文数据量,降低计算成本。
去重的影响及恢复。我们的MT-RAG跟踪研究量化了这一好处,并表明去重仅导致1-3%的准确性下降,这可以通过稍后讨论的技术来恢复。这种微小的损失发生是因为LLM仍然可以通过先前的对话历史访问去重的内容,在避免重复计算的同时保持了质量。
上下文标注补偿准确性扰动。(3) 精心设计的上下文标注可以在很大程度上补偿由上下文块排序和去重引入的微不足道的准确性扰动。如图2c所示,这些标注帮助模型在其内部令牌中重建原始的排序关系,从而减轻准确性损失。评估表明,上下文标注能有效恢复准确性,在多跳推理任务中,甚至能将其提高到超过未排序基线的水平。例如,在NarrativeQA和MultihopRAG上,相对于近似匹配方法,准确性提高了0.3-3.9%(第7.3节),证实了结合上下文标注以增强推理的有效性。需要注意的是,上下文标注不会影响模型的指令遵循能力,因为它只传达了最少的检索元数据,而没有改变用户提示。
ContextPilot的独特性与设计。先前的工作将准确性(如上下文图、智能体记忆)和系统性能(精确前缀缓存)孤立处理。CONTEXTPILOT通过三项贡献独特地弥合了这一差距:(1) 一个带有新颖距离函数的上下文索引,它主动重排文档以最大化前缀重用——将先前系统无法处理的缓存未命中转换成命中;(2) 简洁的上下文标注,使LLM能够在重排后恢复语义优先级;以及(3) 多轮上下文遍历,识别并去重先前已记忆的文档,从而减少预填充开销,尤其是在模型上下文长度受限的情况下。这种协同设计使得效率和质量同时得到提升,这是单一方法无法实现的。
系统架构与接口。CONTEXTPILOT实现了上述三个设计机会,实现了上下文重用且准确性损失可忽略不计。它具有一个清晰、最小化的接口,与常见的检索模块(如FAISS和ElasticSearch)、AI记忆层存储(如Mem0)以及推理引擎(如SGLang和vLLM)兼容。它仅需要在每个引擎的前缀缓存中跟踪请求ID,而不影响现有功能,从而实现快速部署。具体来说,CONTEXTPILOT接收用户提示及其上下文块(检索到的文档、块或记忆),更新上下文以实现有效重用,然后将更新后的上下文传递给推理引擎进行处理。
核心组件。图4展示了CONTEXTPILOT中的关键组件:一个跟踪前缀缓存状态的上下文索引,一个对上下文块进行重排序和调度以最大化缓存命中的上下文排序机制,以及一个移除多轮对话中冗余块的上下文去重机制。以下各节将详细描述每个组件。
设计目标。上下文索引的设计目标是:(1) 高效跟踪推理引擎的前缀缓存状态以实现KV缓存重用;(2) 通过前缀匹配支持快速查找先前存储的KV缓存,从而在存在重叠时实现跨会话的上下文重用;以及(3) 遍历多轮对话中的KV缓存以检测重复的上下文。
索引结构。图5展示了上下文索引的结构示例。左侧面板是索引树,右侧面板是对应的前缀缓存状态。该索引被组织成一棵树,其根节点代表一个空上下文。每个节点对应于前缀缓存中存储的一个前缀,并包含扩展该前缀的子节点。每个节点维护四个属性:(1) 包含上下文块ID的上下文,(2) 从根到该节点的搜索路径,(3) 用于缓存驱逐的访问频率计数器,以及(4) 创建该节点时的聚类距离。
索引创建。索引是通过基于前缀匹配的层次聚类构建的。首先,我们使用上下文之间的重叠率计算它们之间的成对距离。接着,我们迭代地合并最接近的一对,创建一个虚拟节点,其上下文是表示它们共享前缀的排序后交集。最后,每个叶子节点记录其从根开始的搜索路径,从而为跨会话前缀匹配和多轮重复检测实现高效遍历。如图5所示,该过程以叶子节点C1 {2, 1, 3}、C2 {2, 6, 1}和C3 {4, 1, 0}开始。由于C1和C2的距离最小(共享{1, 2}),它们首先合并成一个上下文为{1, 2}的虚拟节点C4。然后C3与C4合并,形成上下文为{1}的根节点C5。最终的树中,C1-C3作为叶子节点存储它们从C5开始的搜索路径,而C4和C5作为虚拟节点代表用于缓存重用的共享前缀。此构建过程的时间复杂度为$O(N^2)$(N为上下文数量),并且可以在CPU和GPU上完全并行化。实际上,为2000个上下文构建索引在CPU上需要8秒,在GPU上需要0.82秒。空间复杂度为$O(N \cdot K)$(K为每个查询的平均上下文块数)。由于索引只存储块ID和元数据而非全文,其空间开销极小。完整的层次聚类伪代码见附录D的算法4。
量化上下文之间的重叠。索引构建中的一个关键挑战是如何量化上下文之间的重叠。我们提出了一个满足两个要求的上下文距离函数:(1) 它能捕捉上下文之间共享文档的数量,以及(2) 它能考虑它们的位置对齐,因为检索系统会根据查询相关性对文档进行排序。为了说明这种设计的必要性,考虑四个上下文:A {3, 5, 1, 7}、B {2, 6, 3, 5}、C {3, 5, 8, 9}和D {2, 6, 4, 0}。一个仅考虑重叠的朴素度量会给A-B、B-C和B-D对分配相同的距离(0.5),因为它们都共享两个文档。然而,B和D在位置1-2共享{2, 6},而A和B在不同位置共享{3, 5}。我们的距离函数(公式1)会给B-D分配一个更小的距离,因为它们的重叠出现在相似的位置,这既反映了重叠的大小也反映了位置的对齐。这种模式无法被传统的距离度量(如余弦、L1或L2相似度)捕捉,因为它们忽略了位置结构。
距离函数定义。我们的距离函数更正式地定义为:
$$d_{ij} = 1 - \frac{|S_{ij}|}{\max(|C_i|, |C_j|)} + \alpha \cdot \frac{\sum_{k \in S_{ij}} |p_i(k) - p_j(k)|}{|S_{ij}|}$$其中,$S_{ij}$表示共享文档的集合,$p_i(k)$是文档k在上下文i中的位置,而$\alpha \in [0.001, 0.01]$确保重叠计数仍然是主导因素,同时加入了位置对齐的考量。
索引更新。上下文索引通过轻量级的请求ID跟踪与推理引擎的前缀缓存保持同步。每个叶子节点都与引擎维护的一个请求ID相关联。当引擎驱逐缓存条目时,它会将相应的请求ID发送给ContextPilot,后者通过一个请求到节点的映射查找受影响的节点并将其移除。空的父节点会被递归地修剪以保持树的紧凑。整个更新成本为$O(h)$,其中h是树的高度,每次驱逐仅需一次遍历。
两种关键操作。上下文索引提供了两种关键操作:
上下文搜索。CONTEXTPILOT会频繁地根据当前上下文搜索先前存储的上下文以实现重用。索引搜索算法(算法1)通过从根节点开始贪婪地向下搜索来高效地定位匹配的上下文,在每一层选择距离最小的子节点,同时记录位置以形成搜索路径。当到达叶子节点或所有子节点距离相等时,搜索停止,这表明找到了最长的共享前缀。更新是局部且高效的:匹配一个内部节点会将新上下文作为其子节点追加($O(1)$),而匹配一个叶子节点则会创建一个新的内部节点,其内容为它们的交集($O(|C|)$)。与K-Means重新聚类或HNSW图重建不同,这些更新不需要重构树,从而实现了动态索引维护且开销极小。
Algorithm 1 Context Index Tree Search
Require: Context C = {b1, . . . , bk}, root node R
Ensure: Search path, best matching node
1: cur ← R, path ← []
2: while cur has children do
3: best ← arg minc∈cur.children, C∩c.blocks≠∅ DIST(C, c)
4: if no overlapping child found then
5: break
6: end if
7: path.APPEND(index of best)
8: if best is leaf then
9: return path, best
10: end if
11: cur ← best
12: end while
13: return path, cur
搜索示例。例如,给定上下文C6 {2, 1, 4},我们在图5的索引中进行搜索。C6首先与根节点的子节点C5比较,发现共享前缀{1},于是下降到C5并记录其位置[0]。在C5处,C6与C4共享{1, 2},但与C3只共享{1},所以它选择C4并追加另一个[0],得到[0,0]。在C4的子节点C1 {1, 2, 3}和C2 {1, 2, 6}处,所有距离都相等,所以搜索停止,并将C4识别为最佳匹配,路径为[0,0]。然后C6被插入到C4的子节点列表中的位置2,形成最终的搜索路径[0,0,2]。
搜索复杂度。搜索复杂度与树的高度成正比。对于具有共同前缀的上下文,$h = O(\log n)$,得到$O(|C| \cdot \log n)$的复杂度,其中n表示存储的上下文数量。根据经验,每次请求的搜索时间约为0.068毫秒(附录C.3),与预填充延迟相比可以忽略不计。
上下文遍历。在多轮对话中,CONTEXTPILOT通过使用存储的搜索路径遍历索引来更新节点上下文长度。从根节点开始,它顺序地沿着路径中的索引前进,直到到达目标节点,然后执行更新。遍历成本为$O(h)$,并且被上述搜索开销所包含。
设计目标。上下文排序机制旨在:(1) 根据当前的前缀缓存状态对检索到的上下文进行重排序,以最大化KV缓存的重用;(2) 在知晓缓存生成和驱逐策略的情况下,将排序后的上下文调度到推理引擎,以提高命中率;以及(3) 插入上下文标注,以恢复排序前的语义并保持准确性。
算法描述。形式上,上下文排序算法(算法2)接收一批请求及其检索到的上下文块作为输入,根据从上下文索引中获得的前缀匹配对它们进行重排序,并返回具有最大化共享前缀的有序上下文。
排序示例。如图6所示,我们从初始化上下文C1 {2, 1, 3}、C2 {2, 6, 1}和C3 {4, 1, 0}开始,随后是新的上下文C6 {2, 1, 4}、C7 {5, 7, 8}和C8 {1, 2, 9}。初始化上下文从其父节点继承前缀(C1、C2从C4继承{1, 2};C3从C5继承{1}),而新上下文则搜索索引(C6和C8匹配C4并继承{1, 2})。每个上下文然后将其匹配的前缀与其余文档按原始顺序连接,产生C1 → {1, 2, 3}、C2 → {1, 2, 6}、C6 → {1, 2, 4}和C8 → {1, 2, 9}。未匹配的上下文(如C7)保持不变并形成独立的-分支。该策略确保了重叠的上下文共享公共前缀,同时保留了非共享文档的排名。
调用时机与开销。该算法在CONTEXTPILOT处理新请求时被调用。其运行时间为$O(|C| \cdot \log n)$,其中n是存储的上下文数量,每个请求大约耗时0.047毫秒(附录C.3),与预填充相比可以忽略不计。
调度策略。在对上下文进行排序后,CONTEXTPILOT必须调度它们的执行以与推理引擎的KV缓存生成和驱逐策略对齐;否则,缓存重用将变得无效。因此,我们设计了一种调度算法,它:(1) 重用在上下文排序期间获得的搜索路径以避免冗余的树查找;(2) 根据搜索路径的第一个元素对上下文进行分组,自然地划分了缓存区域;以及(3) 在每个组内按路径长度降序对上下文进行排序,确保较长的前缀匹配在较短的之前执行。
调度示例。图7说明了这一过程。在基线顺序C6, C3, C7, C8中,有限的缓存容量只允许一个上下文:C6缓存了{1, 2, 4},但C3只重用了{1}并驱逐了{2, 4}。C7导致完全未命中,缓存了{5, 7, 8}并驱逐了所有先前的条目,这接着又迫使C8发生另一次未命中,尽管它与C6共享前缀{1, 2}。这种低效是因为共享前缀的上下文没有被连续执行。我们的调度器将执行顺序重排为C6, C8, C3, C7,将共享前缀的上下文组合在一起。C6首先缓存{1, 2, 4},然后C8在被驱逐前立即重用{1, 2}。C3和C7在之后运行,而不会干扰这种重用,从而最大化了缓存命中率。
调度器开销与对比。我们的调度器对N个上下文执行$O(N)$的分组(按根前缀路径)和$O(N \log N)$的组内排序,实时开销可忽略不计。相比之下,现有的索引方法如RAGCache和SGLang的LPM使用全局前缀选择,在每个决策点重新扫描一个有M个节点的基数树,随着缓存状态的演变,总复杂度为$O(N \log M) + O(N \log N)$。通过顺序地处理分组,我们的方法避免了重复的树搜索,在紧张的KV预算下更好地保持了重用,并使复杂度与M无关。完整的调度伪代码见附录D的算法5。
为什么重排序是安全的。如第3.2节所示,我们对DEmO研究的复现(【Guo et al., 2024a, What makes a good order of examples in in-context learning】)证实,现代LLM对输入顺序的敏感度大大降低(表1),这解释了为什么为提高缓存效率而进行的重排序只引入了微小的准确性扰动,并使得轻量级的校正机制足以应对。
为什么标注仍有帮助。尽管敏感度降低,重排序在某些数据集上仍可能导致轻微的准确性扰动(例如,在QASPER上为-1.1%)。标注通过减少模型对位置信号来推断相关性的依赖来缓解这一问题。在多跳任务中,跨上下文块链接证据需要明确的指导,标注不仅能恢复损失的准确性,还能主动将其提升到超过不重排序基线的水平(例如,在MultihopRAG上使用Qwen3-32B,F1分数+4.0%;见附录C.2)。这种提升在不同模型规模上是一致的:Qwen3-4B在MultihopRAG上提升+1.4%,在NarrativeQA上提升+1.3%;而Qwen3-32B则分别提升+4.0%和+1.2%。注意力图分析(附录B)证实,标注重塑了内部注意力,使其与语义优先级而非位置优先级对齐。
标注机制。我们向LLM提供简洁的标注,指明上下文块的原始相关性排名。重排序上下文会改变这个排名,而这个排名编码了对答案质量至关重要的文档相关性。考虑上下文C6,其中检索器按{2, 1, 4}的顺序返回文档。基线提示是:
[系统提示] → [文档 2] → [文档 1] → [文档 4] → [问题]
重排序与标注示例。为了缓存效率,将其重排序为{1, 2, 4}后,我们在问题前追加一个顺序标注:
[系统提示] → [文档 1] → [文档 2] → [文档 4] → [顺序标注] → [问题]
标注内容。该标注明确指定了原始的相关性优先级:
“请按以下优先顺序阅读上下文:[文档 2] > [文档 1] > [文档 4] 并回答问题。”
效果。这个简短的指令在预填充期间增加了可忽略的令牌开销,但有效地保留了模型根据原始相关性排名关注文档的能力。因此,CONTEXTPILOT在实现积极的缓存优化的同时,准确性扰动很小(在大多数数据集上≤1%),并且在多跳推理任务上通常能提高准确性。
设计目标。上下文去重机制有两个目标:(1) 消除与先前已处理且其KV缓存已存储的上下文块具有完全相同前缀的块,从而最小化冗余的预填充计算;以及(2) 提供标注,告知LLM哪些内容已被去重以及相应信息在先前上下文中的位置。
上下文去重算法。算法3将去重过程形式化。上下文索引为每次对话维护一个记录,包含了先前轮次处理过的所有上下文块,该记录在第一轮的上下文被排序并插入索引时填充(第5节)。给定一个新的上下文,算法会查找这些先前索引的块,识别重叠部分,为每个重复项生成一个位置标注,并注册当前轮次的块以供未来比较。
Algorithm 3 Context De-duplication
Require: New context Cnew, context index I, conversation
ID id
Ensure: Deduplicated context C', location annotations H
1: S ← I.seen_blocks[id] // blocks indexed in prior turns
2: C' ← [], H ← []
3: for each block b in Cnew do
4: if b ∈ S then
5: h ← “refer to [b] in previous conversation”
6: C'.APPEND(h); H.APPEND(h)
7: else
8: C'.APPEND(b)
9: end if
10: end for
11: S ← S ∪ Cnew // register for future turns
12: return C', H
去重示例与开销。我们用一个例子来说明。考虑一个会话,用户C6在第一轮最初检索到上下文{1, 2, 4}。在上下文排序期间(第5节),这些块被插入到上下文索引中,并记录在C6的对话历史中。在第二轮,一个新的查询产生了{1, 5, 2}。算法查询索引的对话记录,识别出{1, 2}已在第一轮缓存,并用位置标注替换它们,只留下新的块{5}需要完全处理。然后,新的块被添加到对话记录中以供未来轮次使用。该算法的运行时间为$O(|C|)$,每个请求大约耗时0.144毫秒(附录C.3),与预填充相比可以忽略不计。
为去重上下文块添加标注。简单地移除重复项可能会降低答案质量。为了保持质量,我们插入了位置标注,引导LLM到对话历史中相应的上下文块。对于C6的第二轮,基线提示是:
[第一轮上下文] → [第一轮问答] → [文档 1] → [文档 5] → [文档 2] → [第二轮问题]
带标注的去重提示。通过去重和位置标注,提示变为:
[第一轮上下文] → [第一轮问答] → [标注 1] → [文档 5] → [标注 2] → [第二轮问题]
标注内容。这里的标注1和标注2是诸如“请参考之前对话中的[文档 1]”和“请参考之前对话中的[文档 2]”之类的引用。这些标注引导LLM到先前的上下文,而无需重复预填充。
评估设置
- 软件实现:CONTEXTPILOT支持SGLang 0.4.6和vLLM 0.10.0,仅需在引擎的前缀缓存中跟踪请求ID,易于集成。
- 基线系统:
- LMCACHE (v0.3.8):代表了提示缓存的SOTA。
- CACHEBLEND:代表了KV缓存匹配的SOTA,与LMCache集成。
- RADIXCACHE:基于SGLang的实现,采用最长前缀匹配(LPM)调度策略。
- HICACHE和RAGCACHE因与RadixCache性能相当或未开源而被排除。
硬件配置:
数据集与模型:
检索与嵌入:
工作负载类型:
多会话RAG
在大型MoE模型上的有效性
- 实验设置:在32个H20 GPU的集群上评估DeepSeek-R1(671B)模型。
- 实验结果(附录A):CONTEXTPILOT在MultihopRAG上将缓存命中率从5%提高到60%,在NarrativeQA上从6%提高到38%。这分别带来了1.81倍和1.52倍的预填充吞吐量提升,并且在扩展到32个GPU时保持了相似的加速比。
多轮RAG
- 实验设置:在MT-RAG数据集上,使用Qwen3-4B、Llama3.1-8B和Qwen3-30B模型,在单个H100 GPU上评估。
- 实验结果(表3a):CONTEXTPILOT通过上下文去重,将首次令牌生成时间(TTFT)相比LMCache降低了3.09-3.45倍,相比RadixCache和CacheBlend也分别有高达2.00倍和1.55倍的加速。
- 准确性:通过位置标注,CONTEXTPILOT保持了准确性,而CacheBlend的准确性下降明显。在某些情况下,CONTEXTPILOT还能提升准确性(例如在Qwen3-4B上从62.56%提升到64.27%)。
多会话、多轮混合RAG
- 实验设置:使用Qwen3-4B模型在H100 GPU上模拟真实部署,并发会话数从2到32不等。
- 实验结果(表3b):在所有并发级别上,CONTEXTPILOT都实现了最低的TTFT。在32个会话时,相比LMCache、RadixCache和CacheBlend,分别实现了2.65倍、1.49倍和1.20倍的加速,证明了其在扩展查询时能有效维持前缀重叠。
多智能体推理
- 实验设置:在Chain-of-Agent (CoA)范式中评估CONTEXTPILOT,使用15个智能体处理MultihopRAG任务。
- 实验结果:通过智能体感知的路由实现KV缓存重用,CONTEXTPILOT在Qwen3-4B上将准确率从48.3%提高到50.2%,吞吐量提高1.8倍;在Llama3.1-8B上将准确率从50.7%提高到54.4%,吞吐量提高2.1倍。
Agentic记忆系统
- 实验设置:将CONTEXTPILOT与流行的AI记忆系统Mem0集成,在LoCoMo基准上使用Qwen3-4B模型进行评估。
- 实验结果:在k=100时,CONTEXTPILOT将TTFT从0.101秒降低到0.055秒(1.83倍加速)。在k=20时,TTFT从0.038秒降低到0.031秒(1.23倍加速),同时准确率从0.440提高到0.460。
排序和调度的贡献(图8)
- 分析:排序和调度两个组件都对缓存命中率有增量贡献。在SGLang和Qwen3-32B上,仅排序将命中率从8.49%提升至20.56%,加上调度后达到33.97%,总体提升4倍。
长时运行工作负载下的性能(附录C.1)
- 分析:在整个工作负载期间,CONTEXTPILOT持续维持约34%的缓存命中率,而基线约为7%,表现出稳定的5倍优势(图12)。累计缓存令牌数是基线的4.27倍(图13),证明了持续的预填充加速效果。
上下文重用下的模型准确性(附录C.2)
- 分析:单独重排序导致的F1分数变化≤1%。添加标注不仅弥补了这一差距,甚至将准确性提高了1.4-4.4%(表5)。
上下文索引构建开销(表3c)
- 分析:索引构建开销轻量且可扩展。在A6000 GPU上,处理12K个上下文耗时7.48秒。每个请求的开销约为0.2毫秒(附录C.3),与秒级的预填充延迟相比可以忽略不计。
系统开销(零上下文重叠)
- 分析:在没有检索重叠的最坏情况下,CONTEXTPILOT对一个1K上下文的一小时作业仅增加了0.72秒的预填充延迟,表明其运行时开销极小。
随上下文长度增长的性能(图9)
- 分析:在不同的检索深度k(3, 5, 10, 15)下,CONTEXTPILOT始终保持最高的预填充吞吐量。随着k的增加,基线吞吐量下降,而CONTEXTPILOT保持稳定,表现出1.3-2.0倍的持续加速。
上下文标注的有效性(附录B)
- 分析:注意力热图分析证实,标注能有效引导模型关注语义上相关的文档,而不是物理顺序上的文档,从而提高下游任务的准确性。
前缀缓存大小的影响
- 分析:CONTEXTPILOT能从更大的缓存中获益更多。在从A6000(48GB)切换到H100(80GB)时,其缓存命中率提升显著(例如SGLang从29.64%提升到33.97%),因为其上下文排序策略能系统地利用额外的缓存容量。
本文介绍了CONTEXTPILOT,一个通过上下文重用来加速长上下文预填充的系统,且准确性损失可忽略不计。CONTEXTPILOT将外部输入统一表示为上下文块,并应用一套通用的机制——上下文索引、排序、去重和简洁的标注——来提升跨多样化工作负载的前缀KV缓存重用率。其模块化设计为上下文工程、管理和优化的新系统及算法研究提供了可能。展望未来,我们设想CONTEXTPILOT将成为一个协同优化框架,共同决定输入什么上下文以及如何服务它,从而最大化推理效率和输出质量。
端到端结果。表4展示了DeepSeek-R1在MultihopRAG和NarrativeQA数据集上的端到端结果,评估在16个H20和32个H20 GPU上进行。CONTEXTPILOT在MultihopRAG和NarrativeQA上分别实现了1.81倍和1.52倍的吞吐量提升,这些加速效果在使用上下文感知路由的16和32个GPU部署中保持一致。
分析动机。由于上下文工程和上下文学习对模型推理有很大影响,我们分析了在引入明确标注以回忆原始相关性排名时模型的注意力模式。
分析结果。图10和图11比较了Qwen3和LLaMA3.3在这种设置下最后一层注意力图。当给予明确的文档优先级线索时,尽管架构不同,两个模型都表现出一致的注意力行为。它们正确地关注文档令牌([Doc 1], [Doc 2], [Doc 3]),这反映了它们意识到重排序序列与原始序列之间的不匹配,如标注区域的查询与上下文区域的键之间的交集所示。随着上下文与原始序列重新对齐,两个模型在解析([Doc 1])和([Doc 3])时都强调了([Doc 2]),表明线索([Doc 2] > [Doc 1] > [Doc 3])有效地引导了跨文档的注意力。因此,明确的标注重塑了内部注意力,使其与语义优先级而非位置优先级对齐。这一发现支持了ContextPilot的一个核心假设:明确的标注可以通过重新建立与原始检索语义的对齐,来弥补因重排序和过滤而损失的准确性。
缓存命中率随时间变化。图12显示了缓存命中率随工作负载进展的演变。ContextPilot在整个工作负载期间保持约34%的缓存命中率,而基线为7%,显示出持续的5倍改进,这并非暂时的预热效应。
累积缓存令牌数。图13将累积缓存令牌数作为基数树前缀重用的度量。在Llama3.3-70B-Instruct上,ContextPilot在完成时实现了1033万缓存令牌,而基线为242万(4.27倍);在Qwen3-32B上为1050万对275万(3.82倍)。“无调度”变体(685万,2.83倍)证实了排序和调度都对改进做出了贡献。
各组件对准确性的贡献。表5提供了ContextPilot各组件对准确性贡献的详细分解。
各组件开销。表6报告了ContextPilot组件的每请求开销,测量于NVIDIA A6000上2K个请求,k = 15。
索引构建与请求调度。算法4提供了通过层次聚类构建上下文索引的完整伪代码,算法5详细说明了基于树的请求分组和调度过程。
Algorithm 4 Context Index Construction via Hierarchical
Clustering
Input: Batch of n contexts S = {s1, s2, . . . , sn}, distance
function d(·, ·)
1: // Phase 1: Distance computation and clustering
2: D ← pairwise distances using d(si, sj ) // GPU or CPU
3: Z ← LINKAGE(D) // hierarchical clustering
4: // Phase 2: Build tree with deduplication
5: for i = 1 to n do
6: Create leaf node vi with vi.content ← si
7: if ∃ vj with vj.content = si then
8: Redirect vi → vj ; vj .freq += 1 // dedup
9: end if
10: end for
11: for each merge (c1, c2, δ) in Z do
12: v ← new internal node
13: v.content ← c1.content ∩ c2.content
14: v.children ← [c1, c2]; c1.parent, c2.parent ← v
15: end for
16: Remove empty internal nodes; relink children to grand-
parents
17: // Phase 3: Top-down prefix-aligned reordering
18: Compute search paths from root to all nodes
19: for each node v in BFS order from root do
20: if v is not root and v.parent.docs ≠ ∅ then
21: v.docs ← v.parent.docs ⊕ (v.docs \ v.parent.docs)
22: end if
23: end for
24: return reordered contexts from leaf nodes, search paths
Algorithm 5 Search-Path-Based Request Grouping and Scheduling
1: Input: N contexts with search paths {P1, . . . , PN }
from ordering
2: Output: Scheduled execution order
3: // Phase 1: Group by root prefix // O(N)
4: G ← empty map
5: for i = 1 to N do
6: key ← Pi[0] // first element of search path
7: G[key].APPEND(i)
8: end for
9: // Phase 2: Sort within each group // O(N log N)
10: for each group G in G do
11: Sort G by |Pi| descending // longer paths first
12: end for
13: // Phase 3: Order groups and flatten
14: Sort groups by |G| descending // largest groups first
15: return concatenation of all groups