Learned Prefix Caching for Efficient LLM Inference
Learned Prefix Caching for Efficient LLM Inference
作者/机构: Dongsheng Yang, Austin Li, Kai Li, Wyatt Lloyd (Department of Computer Science, Princeton University)
A1 主要贡献
核心问题:前缀缓存(Prefix caching)是降低大语言模型(LLM)推理成本的关键技术,但当前普遍使用的最近最少使用(LRU)驱逐算法与理想的最优算法之间存在巨大性能差距。现有系统(如 vLLM 和 SGLang)采用 LRU 策略,它根据最近的访问时间来驱逐 KV 块,这种策略对于 LLM 对话工作负载而言并非最优,因为用户在多轮对话中可能会有自然的停顿,导致有用的前缀因长时间未被访问而被错误地驱逐。如图 1 所示,LRU 与拥有完美未来知识的 Oracle 策略在缓存命中率上存在显著差距。
研究目标:为了弥补这一差距,本文旨在设计一种专门针对 LLM 对话系统的轻量级学习型前缀缓存框架(LPC),以解决两个核心挑战:1)通过准确预测对话继续的可能性,实现远高于 LRU 的缓存命中率;2)确保学习组件的内存和计算开销可以忽略不计,从而在资源受限的 LLM 服务环境中实现更低的推理延迟和更高的吞吐量。
创新点:
1. 首个学习型前缀缓存方法:本文提出了首个利用机器学习来优化 LLM 推理效率的学习型前缀缓存方法 LPC,通过提高前缀的重用率来提升性能。
2. 轻量级对话继续概率预测:设计了一种通过分析对话内容来学习和预测对话继续概率的方法。该方法结合了文本嵌入和一个小型的多层感知机(MLP),能够在内存和计算开销极小的情况下,为缓存驱逐提供预测性指导。
3. 显著的性能提升:通过在三个真实世界数据集上的广泛评估,LPC 实现了以下成果:
* 在达到相同命中率的情况下,所需缓存大小减少了 18–47%。
* 在一个模拟的解耦服务环境中,LLM 预填充(prefilling)吞吐量提高了 11%。
* 高达 7% 的请求其首个令牌(first token)延迟降低了 42–75%。
- 系统级正交性:LPC 作为一种系统级的缓存管理策略,与其他优化技术(如 FlashAttention 或推测解码)是正交的,可以结合使用以获得复合效益。
图 1: LRU 与拥有完美未来对话继续知识的理想化 Oracle 策略在命中率上存在巨大差距(LMSys 数据集)。
A3 背景知识
LLM 推理中的 KV 缓存。KV 缓存存储键值(Key-Value, KV)状态,以支持自回归 LLM 推理。这些 KV 状态是每个 token 上下文信息的计算表示。KV 缓存在 LLM 推理的两个阶段中被大量使用。首先,在预填充(prefilling)阶段,系统处理整个输入提示。对于提示中的每个 token(例如 tokeni),其对应的 KV 状态 kvi 会被计算并存储在 KV 缓存中。此阶段以生成第一个输出 token 结束。随后,在解码(decoding)阶段,系统利用 KV 缓存中该序列所有先前 token 的 KV 状态(代表了迄今为止的全部上下文)来生成下一个 token。为了高效管理内存并减少碎片,KV 缓存通常将这些状态组织成固定大小的“KV 块”,块大小通常约为 16 个 token【10,Efficient memory management for large language model serving with pagedattention by Woosuk Kwon et al., 2023, Proceedings of the 29th ACM Symposium on Operating Systems Principles (SOSP)】。当一个序列完成后,与其相关的 KV 状态通常会从 KV 缓存中清除。因此,这个 KV 缓存的主要作用是在单个生成序列内消除对 token 上下文的重复计算。
图 2: 在 LLM 推理中使用 KV 缓存和前缀缓存的示例。
前缀缓存。标准的 KV 缓存通常在请求完成后丢弃 KV 状态。这妨碍了在相关请求间的重用,例如继续先前对话的请求,从而迫使系统对共享的上下文进行代价高昂的重新预填充。前缀缓存通过保留已完成序列的 KV 状态来解决此问题。当新请求到达时,系统会查询前缀缓存以寻找与输入最长匹配的前缀,例如 token1, . . . , tokenm。如果找到了其 KV 状态(kv1, . . . , kvm),这些状态就会被提供给当前请求的 KV 缓存。这样,预填充仅需对输入中剩余的、未匹配的部分(从 tokenm+1 开始)进行,从而显著减少计算量。图 2 展示了一个例子,其中 "What can" 的预填充被跳过了,因为它在前缀缓存中匹配成功。
前缀缓存驱逐。前缀缓存的容量有限,因为它使用的是宝贵的 GPU 内存。当前缀缓存变满或 KV 缓存需要更多内存时,系统必须驱逐现有的 KV 块。选择驱逐哪些块至关重要:一个精心设计的驱逐策略能直接提高前缀缓存的命中率——即未来请求找到匹配前缀的比例——从而减少冗余计算,并提高吞吐量和延迟。
现有 LRU 策略的局限性。然而,现有系统使用 LRU 策略,该策略驱逐最近最少使用的块。这种方法在传统缓存领域相当有效,但在 LLM 服务环境中表现不佳。前缀缓存工作负载表现出不同的时间和空间局部性模式。在人与 LLM 的交互中,用户在不同轮次之间会自然地停顿下来阅读、思考或撰写后续查询,这意味着有用的前缀可能会在很长一段时间内未被触及,尽管它们仍然高度相关。此外,LRU 无法捕捉对话的语义或任务相关上下文——这些因素通常比最近使用情况更能预测未来的重用。一个围绕复杂、未解决任务的对话比一个简短、封闭式交流的对话更有可能继续。这些差异促使我们设计新的驱逐策略,这些策略要超越最近性,并考虑 LLM 工作负载中独特的行为和语义模式。
A2 方法细节
本节介绍了一种学习型前缀缓存方法,旨在显著提升 LLM 推理中缓存管理的技术水平。我们的目标是通过智能地选择在需要空间时驱逐哪些键值(KV)块来提高前缀缓存的效率,从而降低预填充延迟并增加吞吐量。为实现此目标,我们的方法必须满足三个关键要求:
* 高命中率:该方法应在各种工作负载下,缓存命中率始终显著优于 LRU 基线。
* 改善延迟和吞吐量:该方法必须高效,不引入瓶颈,从而实现更低的延迟和更高的吞吞吐量。
* 最小化开销:系统在 GPU 内存使用和计算成本方面应施加可忽略的开销。
3.1 LPC 框架概述
图 3: LPC 概览。LPC 框架包含一个预测器和一个驱逐算法。
LPC 框架设计。我们提出了学习型前缀缓存(LPC)来满足上述要求。图 3 展示了 LPC 的设计。该方法与普通前缀缓存的主要区别在于,LPC 使用一个预测器来告知替换算法应驱逐哪些块,而不是选择最近最少使用的块。
LPC 预测器工作流程。LPC 预测器接收对话历史作为输入,并输出一个概率分数,该分数表示对话继续的可能性。其步骤如下:
1. 数据解析:从对话历史中提取用户提示。
2. 文本嵌入:从解析出的用户提示生成一个语义嵌入向量。
3. 基于 MLP 的分类:将文本嵌入和其他特征输入一个多层感知机(MLP)分类器,以预测对话继续的概率 p。
LPC 驱逐算法组件。LPC 驱逐算法有两个运行组件:
1. 缓存插入与驱逐:根据更新后的概率 p′ 将 KV 块插入到一个最小堆(min-heap)中;当前缀缓存满时,驱逐 p′ 最小的 KV 块。
2. 概率更新:周期性地对预测的继续概率 p 应用一个衰减函数,以获得一个考虑了已流逝时间的更新概率 p′。
3.2 LPC 预测器
本节描述了 LPC 预测器的三个组成部分:数据解析、文本嵌入模型和 MLP 分类器。
数据解析。为了向模型提供足够的上下文以进行准确预测,我们从对话中提取文本和数值特征。这些特征包括新用户提示的文本以及之前 N 轮对话的用户提示。我们的实现中使用 N = 4,以便在保持较低资源使用的同时,尽可能多地保留上下文信息。出于类似原因,我们不包括模型的响应,因为它们比用户提示包含的新信息要少,这一点在附录的消融研究中得到了验证。此外,对话长度(表示为已发生的轮次数)作为一个标量特征被包含在内,这对于超过 N 轮的对话尤其具有信息价值。
文本嵌入模型。为了将原始文本转换为有意义的向量表示,我们使用了 sentence transformer 模型家族中的 multilingual-e5-small 模型【27,Multilingual e5 text embeddings: A technical report by Liang Wang et al., 2024, arXiv preprint arXiv:2402.05672】(采用 MIT 许可证)。我们选择该模型是因为它能够表示重要模式且效率很高。该模型有 1.18 亿参数。这类高效模型通常对输入长度有限制;对于 multilingual-e5-small,长度限制为 512 个 token。因此,一个关键挑战是在压缩超过长度限制的输入的同时保留基本信息。我们采用一种策略性的 token 分配方法:将 512 个 token 的预算平均分配给 N+1 个用户请求轮次中的每一轮。对于每一轮分配的 token 预算,一半的 token 来自提示的开头,另一半来自结尾。该策略的依据是研究表明,文本输入中最重要的信息通常位于这些极端位置【11,Lost in the middle: How language models use long contexts by Nelson F Liu et al., 2023, arXiv preprint arXiv:2307.03172】。文本嵌入模型的输出是一个 384 维的向量。
基于 MLP 的分类模型。分类模型是一个 3 层的 MLP,每层有 128 个隐藏神经元。我们选择它是因为它非常高效,同时具备足够的学习能力来完成我们的任务。它接受一个拼接的特征向量作为输入,该向量由 384 维的文本嵌入向量和代表对话轮次数的标量值组成。模型输出一个介于 0 到 1 之间的单一概率值,表示对话在当前轮次后继续的置信度。
开销最小化设计。我们做出了几项设计决策以最小化预测器的开销。首先,我们使用小型的文本嵌入和 MLP 模型,它们的 GPU 内存开销很低。其次,预测器对每个用户请求只运行一次,而请求速率通常低于每 GPU 每秒 10 次。第三,预测器是一个独立组件,与 LLM 推理并行运行。最后,一个预测器可以与多个 LLM 推理实例配对,通过批处理进一步分摊内存和计算开销。
3.3 训练
本节概述了 LPC 的训练方式,包括训练数据收集、训练配置和再训练。
训练数据。训练和验证数据从目标数据集(如第 4.2 节所述)的专用对话分区中收集。该分区被严格隔离,不用于任何在线评估的数据集。分区的大小是数据集的一半。为每个数据集训练一个单独的模型。
训练配置。训练过程如下。首先,在训练期间,只更新 MLP 分类器的权重;预训练的文本嵌入模型保持冻结。这种部分微调方法允许 MLP 专门适应对话继续预测任务,同时确保框架保持轻量级。对于频繁的再训练,完全微调在操作上成本高昂,而只训练微小的 MLP 头部则非常快(通常少于 10 分钟),使得每日自适应成为可能。其次,训练数据集的文本嵌入在第一个 epoch 后被预计算并缓存,以显著加速训练。后续的 epoch 直接重用这些缓存的嵌入,因为文本嵌入模型是冻结的,从而绕过了重复计算,节省了 90% 的训练时间。采用了二元交叉熵损失函数。为了处理类别不平衡问题,损失函数对少数类别赋予了更高的权重。优化器是 Adam,学习率为 $5 \times 10^{-4}$。训练会持续到收敛(在我们的评估中为 20 个 epoch 以内),并在验证数据集上损失最低的检查点被保存下来用于在线推理。
再训练。预测器是在历史数据上离线训练的。为了适应用户行为的逐渐变化或不断演变的对话动态,并减轻因分布外(OOD)数据导致性能下降的风险,周期性的离线再训练是主要策略。我们建议每日进行再训练,这是基于机器学习的缓存策略为适应不断变化的查询模式而采取的标准做法【23,{HALP}: Heuristic aided learned preference eviction policy for {YouTube} content delivery network by Zhenyu Song et al., 2023, 20th USENIX Symposium on Networked Systems Design and Implementation (NSDI 23)】。训练大约需要 10 分钟,使得这个计划是可行的。对于 OOD 风险极高的用例,例如由快速变化的当前事件驱动的应用,系统可以利用在线学习进行增量更新。然而,对于典型的对话工作负载,对话继续的一般模式相对稳定,使得频繁的离线再训练成为一个充分且操作上更简单的选择。
3.4 驱逐算法
本节更详细地介绍 LPC 驱逐算法。我们的设计在很大程度上遵循了 LRU 前缀缓存驱逐算法,但使用动态的对话继续概率而非最后访问时间来对对象进行排序。
3.4.1 数据结构与操作
核心数据结构与操作。关键数据结构是一个最小堆(min-heap)。堆中的块根据其关联对话的动态计算的继续概率进行优先级排序。与每个序列的 KV 块相关联的主要输入是其初始预测的继续概率和其最后访问的时间戳。
操作包括:
* 插入:当一个序列的解码完成时,其 KV 块会从 KV 缓存中移除,并根据其初始预测的继续概率插入到前缀缓存的最小堆中。
* 提升:如果一个传入用户请求的前缀与前缀缓存中的一系列 KV 块匹配,这些块被视为“命中”。然后它们被提升到 KV 缓存中,并因此从前缀缓存的最小堆中移除。
* 概率更新:该算法执行周期性的概率衰减,以便能够驱逐那些继续概率被高估的 KV 块。在可配置的固定时间间隔(默认为 10 秒),LPC 会遍历前缀缓存中的所有 KV 块。对于每个块,其关联的继续概率会使用下面的衰减函数重新评估。然后,最小堆会重新组织以准确反映这些更新后的概率。
* 驱逐:当前缀缓存达到其容量,或需要回收内存时(例如,由于 KV 缓存的扩展),驱逐过程将被触发。位于最小堆根部的 KV 块被选中并从前缀缓存中移除。
3.4.2 概率衰减函数
处理高估概率。为了处理那些因继续概率预测被高估而可能永久停留在缓存中的 KV 块,我们周期性地更新它们的概率。我们实现了一个衰减函数,随着自上次交互以来时间的流逝来调整对话的继续概率。该函数的输入包括初始预测的继续概率 ($prob_{original}$),当前系统时间 ($time_{cur}$),以及该 KV 块的最后访问时间戳 ($time_{last}$)。其输出是当前经时间调整后的概率 ($prob_{cur}$),计算方式如下:
$$prob_{\text{cur}} = \frac{prob_{\text{original}} \times decay}{prob_{\text{original}} \times decay + (1 - prob_{\text{original}})}$$衰减因子。衰减因子本身是流逝时间的指数函数:
$$decay = \exp(-(time_{\text{cur}} - time_{\text{last}}) \times scale)$$在这里,scale 是一个控制衰减率的超参数。它是通过经验调优的,一个好的经验法则是将其设置为平均轮次间隔的倒数。由于在我们的数据集中这个间隔大约是 100 秒,我们在所有实验中设置 scale = $10^{-2}$。这个机制确保了来自不活跃对话的前缀会获得逐渐降低的继续概率,使它们更有可能成为被驱逐的候选对象。
3.5 处理由多个对话共享的 KV 块
共享块的挑战。前缀缓存中的一个重要挑战来自于被多个不同对话共享的 KV 块,例如那些由常见的系统提示或相同的初始用户查询产生的块。如果元数据管理简单地用最新的分数覆盖这些块的分数,可能会导致次优的驱逐决策,因为一个被多个对话共享的热门前缀,如果最近相关的对话得分较低,就可能被错误地驱逐。
LPC 的处理逻辑。我们的算法为这些共享块加入了特定的逻辑。共享块的 $time_{last}$ 总是更新为其最近的访问时间戳。此外,当一个共享块作为新对话或现有对话的一部分被访问时,其用于驱逐决策的有效继续概率会根据传入的元数据计算得出。这种在所有共享该块的对话中进行的最大池化操作表示为:
$$prob_{\text{cur}}^{\text{block}} = \max(prob_{\text{cur}}^{\text{block}}, prob_{\text{cur}}^{\text{new}})$$这个策略确保了一个被多个对话共享的块,只要其中至少有一个对话极有可能继续并且最近是活跃的,就不会因为与一个不太可能继续的对话相关联而被过早驱逐。
A4 实验环境
- 模型架构:使用 Qwen3-32B-FP8 模型【18,Qwen3 by Qwen3 Team, 2025, https://github.com/QwenLM/Qwen3】进行推理,这是一个拥 有 320 亿参数的推理模型,采用 Apache 2.0 许可证。
-
硬件配置:
- GPU:NVIDIA H100 GPU,配备 80 GB HBM3 内存。所有实验均在单 GPU 上运行。
- CPU:8 核 CPU。
- 内存:64 GB CPU 内存。
- 缓存配置:考虑到模型和 PyTorch 本身占用的内存,可用于 KV 缓存和前缀缓存的总内存为 40 GB。为 LPC 的预测器预留 1 GB GPU 内存(约 560 MB 用于 CUDA 上下文,240 MB 用于 float16 精度的模型权重,200 MB 用于运行时内存)。在与 LRU 的公平比较中,LPC 的可用前缀缓存大小会减去这 1 GB。
-
软件配置:
- 实现:在 vLLM 服务框架【10,Efficient memory management for large language model serving with pagedattention by Woosuk Kwon et al., 2023, Proceedings of the 29th ACM Symposium on Operating Systems Principles (SOSP)】(Apache-2.0 许可证)之上实现 LPC 原型。对话继续预测器作为外部组件运行,而缓存驱逐策略直接集成到 vLLM 的前缀缓存中。
-
数据集:
- LMSys【29,Lmsys-chat-1m: A large-scale real-world llm conversation dataset by Lianmin Zheng et al., 2023a, arXiv preprint arXiv:2309.11998】:大规模用户与聊天机器人的对话集合,拥有有限使用许可。
- ShareGPT【20,Sharegpt dataset by ShareGPT Team, 2023, https://huggingface.co/datasets/anon8231489123/ ShareGPT_Vicuna_unfiltered】:与 ChatGPT 的对话,采用 Apache-2.0 许可证。
- Chatbot-Arena【2,Chatbot arena: An open platform for evaluating llms by human preference by Wei-Lin Chiang et al., 2024, Forty-first International Conference on Machine Learning】:用户比较各种 LLM 的平台,采用 CC 许可证。
- 数据集统计:见下表 1。
Table 1: 数据集属性。
-
工作负载生成:由于公共数据集缺乏时间戳,本文采用两阶段策略模拟真实负载。
- 预运行阶段:从均值为 $\lambda_{chat}$ 的指数分布中采样对话内的聊天间隔,并间隔对话开始时间以使并发活跃对话数低于目标阈值 $N_{conv}$(默认为 200)。
- 运行时阶段:每个对话在其指定时间开始,后续请求在前一个 LLM 响应的实际完成时间上加上采样的用户间隔后发生。评估运行 20 分钟。
A4 实验结果
命中率与缓存大小的关系(4.3)
- 实验内容:比较 LPC 和 LRU 在不同总 token 容量(从 60K 到 160K tokens,约 15.6 GB 到 40 GB)下的前缀缓存命中率。聊天间隔参数 $\lambda_{chat}$ 固定为 100。
- 实验结果:如图 4 所示,LPC 在所有数据集和缓存大小上均持续优于 LRU。相对 LRU 的提升在 LMSys 上为 13% 至 38%,在 ShareGPT 上为 14% 至 98%,在 Chatbot-Arena 上为 15% 至 30%。
- 分析结论:为达到与 LRU 相同的命中率,LPC 所需的内存显著减少。例如,在 Chatbot-Arena 数据集上,LPC 用 85K tokens 即可达到 30% 的命中率,而 LRU 需要 150K tokens。这表明 LPC 能将 GPU 内存节省 7.2–18.8 GB(在 40 GB 缓存中),从而大幅提升了前缀缓存的内存效率。
图 4: 命中率 vs. 缓存大小。LPC 始终比 LRU 获得更高的命中率,这转化为在达到与 LRU 相同命中率时,所需缓存大小减少 18–47%。
聊天间隔对命中率的影响(4.4)
* 实验内容:固定缓存大小为 40 GB,改变聊天间隔参数 $\lambda_{chat}$ 来模拟不同的用户思考时间,评估 LPC 的鲁棒性。
* 实验结果:如图 5 所示,LPC 在所有数据集和间隔设置下均持续优于 LRU,实现了 5% 到 18% 的相对命中率提升。
* 分析结论:LPC 能够很好地适应不同的用户行为和对话节奏,在多样化的对话工作负载中保持强大的性能。
图 5: 命中率 vs. 聊天间隔。LPC 的命中率始终高于 LRU,相对提升 5–18%。
吞吐量和延迟(4.5)
* 实验内容:使用微基准测试模拟一个仅处理预填充(prefilling-only)的解耦推理服务器,以评估命中率提升对吞吐量和延迟的影响。工作负载为 1000-token 上下文加 50-token 新提示的合成负载。
* 实验结果:
* 吞吐量:如图 6 所示,命中率从 35% 提升到 42%(LPC 在图 5a 中实现的提升),预填充吞吐量从 10.01 req/s 增加到 11.11 req/s,提升了 11.10%。
* 延迟:如图 7 所示,上下文缓存命中可以显著降低首个令牌时间(TTFT),与未命中相比最多可减少 73%。7% 的命中率提升意味着 7% 的请求其 TTFT 可以减少 42-75%。
- 分析结论:LPC 带来的更高命中率能有效转化为更高的系统吞吐量和更低的用户感知延迟。
图 6: 预填充吞吐量 (req/s) 作为命中率的函数。
图 7: 在不同上下文和提示长度下,上下文缓存未命中与命中的首个令牌时间(TTFT)比较。
内存和计算开销(4.6)
* 实验内容:评估 LPC 预测器引入的开销。
* 实验结果:
* 内存开销:为 LPC 的预测器预留了 1 GB GPU 内存,占 H100 GPU 总内存的 1.25%。
* 计算开销:LPC 的 GPU 计算开销极小,因为其高效的预测器模型每 GPU 每秒可处理多达 1000 次预测。
- 分析结论:LPC 引入的开销微不足道,其带来的性能收益远超其资源成本。附录中的详细分析进一步证实了这一点。
A7 补充细节:相关工作
前缀缓存与驱逐策略。许多工作探索了通过前缀缓存优化 LLM 推理,但没有将重点放在驱逐策略如何适应 LLM 工作负载上。PromptCache【5,Prompt cache: Modular attention reuse for low-latency inference by In Gim et al., 2024, Proceedings of Machine Learning and Systems】通过预计算的注意力状态重用常见的提示段落。RAGCache【9,Ragcache: Efficient knowledge caching for retrieval-augmented generation by Chao Jin et al., 2024, arXiv preprint arXiv:2404.12457】为 RAG 提出了前缀感知的 Greedy Dual-Size Frequency 驱逐策略,但并非针对对话工作负载。CacheBlend【29,Cacheblend: Fast large language model serving for rag with cached knowledge fusion by Jiayi Yao et al., 2025, Proceedings of the Twentieth European Conference on Computer Systems】使用前缀缓存和交叉... SGLang【30,Sglang: Efficient execution of structured language model programs by Lianmin Zheng et al., 2024, Advances in Neural Information Processing Systems】构建了 RadixAttention,将 KV 序列缓存在基数树中以实现高效的前缀匹配,而 vLLM【10,Efficient memory management for large language model serving with pagedattention by Woosuk Kwon et al., 2023, Proceedings of the 29th ACM Symposium on Operating Systems Principles (SOSP)】将逻辑 KV 块映射到全局物理缓存。两者都采用了最近最少使用(LRU)的驱逐策略。Dynamo【12,Nvidia dynamo by NVIDIA Research Team, 2025, https://github.com/ai-dynamo/dynamo】使用先进先出(FIFO)驱逐, 而 CachedAttention【4,Cost-efficient large language model serving for multi-turn conversations with cachedattention by Bin Gao et al., 2024, Proceedings of the 2024 USENIX Annual Technical Conference (USENIX ATC 24)】提出了一种基于传入请求队列的、感知调度器的驱逐方案,用于离线批处理推理。Marconi【13,Marconi: Prefix caching for the era of hybrid llms by Rui Pan et al., 2025, Proceedings of the 7th Conference on Machine Learning and Systems (MLSys 2025)】在 LRU 驱逐中考虑了计算节省,但该方法专为状态空间模型定制,不适用于 transformer 架构。
基于机器学习的缓存算法。现有的缓存研究探索了使用机器学习来逼近最优驱逐决策。LeCaR【26,Driving cache replacement with {ML-based}{LeCaR} by Giuseppe Vietri et al., 2018, 10th USENIX Workshop on Hot Topics in Storage and File Systems (HotStorage 18)】和 CACHEUS【19,Learning cache replacement with {CACHEUS} by Liana V Rodriguez et al., 2021, 19th USENIX Conference on File and Storage Technologies (FAST 21)】应用在线学习和专家集成框架来进一步学习缓存驱逐策略。LRB【22,Learning relaxed belady for content distribution network caching by Zhenyu Song et al., 2020, Proceedings of the 17th USENIX Symposium on Networked Systems Design and Implementation (NSDI 2020)】、MAT【28,A learned cache eviction framework with minimal overhead by Dongsheng Yang et al., 2023, arXiv preprint arXiv:2301.11886】和 HALP【23,{HALP}: Heuristic aided learned preference eviction policy for {YouTube} content delivery network by Zhenyu Song et al., 2023, 20th USENIX Symposium on Networked Systems Design and Implementation (NSDI 23)】在内容分发网络中利用机器学习,展示了相比最佳启发式算法和 LRU 实现了 10% 到 20% 的命中率提升和成本节省。然而,由于缓存范式的巨大差异,这些算法不能直接应用于 LLM 前缀缓存。
用于 LLM 推理的学习型优化。最近的 LLM 推理系统使用学习方法来提升性能。【31,Response length perception and sequence scheduling: An llm-empowered llm inference pipeline by Zangwei Zheng et al., 2023b, https://arxiv.org/abs/2305.13144】添加了预测提示长度的指令,以便将具有相似延迟的请求分组,从而相比常规批处理推理提高了吞吐量 。Fu 等人【3,Efficient llm scheduling by learning to rank by Yichao Fu et al., 2024, https://arxiv.org/abs/2408.15792】 和 Qiu 等人【16,Power-aware deep learning model serving with -Serve by Haoran Qiu et al., 2024a, 2024 USENIX Annual Technical Conference (USENIX ATC 24)】【17,Efficient interactive llm serving with proxy model-based sequence length prediction by Haoran Qiu et al., 2024b, https://arxiv.org/abs/ 2404.08509】训练了更小的独立模型来预测输出长度,发现最短作业优先的批处理相比先到先服务显著降低了每 token 的延迟。最后,Jain 等人【8,Intelligent router for llm workloads: Improving performance through workload-aware load balancing by Kunal Jain et al., 2025, https://arxiv.org/abs/2408.13510】探索了一种基于启发式引导强化学习的路由器,用于实现工作负载感知的调度。我们 的 LPC 工作通过将其学习型优化聚焦于 LLM 推理的一个不同组件:前缀缓存,来补充这项工作。
A5 结论
本文提出了首个学习型前缀缓存 LPC,旨在提高 LLM 推理效率。它利用机器学习对历史数据进行分析,预测对话继续的概率,从而提升了缓存命中率,优于基于 LRU 的方法。
我们的评估表明,LPC 在所有三种工作负载中均持续优于基线,实现了高达 43% 的缓存大小利用率降低。更高的命中率转化为显著降低的延迟和 11% 的预填充吞吐量提升。这是朝着 LLM 服务系统中智能和自适应内存管理迈出的一步。
未来工作:尽管我们的方法相比 LRU 有显著改进,本研究仍有几个局限性计划在未来工作中解决。首先,我们的评估依赖于带有合成时间戳的工作负载,可能无法完全捕捉真实世界的用户行为。其次,我们的比较主要针对 LRU;与更广泛的启发式驱逐策略进行评估将提供更全面的视角。第三,虽然 LPC 的原理是模型无关的,但我们的实验集中于单一模型家族;在不同架构(如 Llama, Mistral)上展示其有效性将进一步加强我们关于其普适性的主张。这项工作也为未来的研究开辟了几个有希望的方向,例如对特定领域(编码、代理任务)进行更细粒度的命中率分析,以及将 LPC 扩展以处理多模态输入。
A6 附录
A 附加实验
A.1 LPC 计算开销
评估指标与设置。本节扩展了第 4.6 节,以解决我们提出的方法 LPC 的计算开销问题。我们使用 NVIDIA 数据中心 GPU 管理器收集的三个细粒度指标来评估 GPU 计算活动:SM Active、SM Occupancy 和 Tensor Pipe Active。实验设置包括三种配置:我们提出的 LPC 方法、基线 LRU 方法,以及 LRU w/ predictor 配置(该配置将 LPC 预测器与标准 LRU 策略集成,用于测量预测器固有的计算开销)。实验在 LMSys 数据集上进行,缓存大小为 40 GB,聊天间隔参数 $\lambda_{chat}$ 为 100。
Table 2: 在 LMSys 上,缓存大小为 40 GB 时的 GPU 计算活动。
结果与分析。表 2 的结果显示,LPC 使整体 GPU 计算量减少了约 5%。具体来说,LPC 在 SM Active 上降低了 4.89%,在 SM Occupancy 上降低了 5.22%,在 Tensor Pipe Active 上降低了 5.23%,这凸显了其由于更高的前缀缓存命中率而有效减少了整体计算负载。其次,为了评估 LPC 预测器引入的计算开销,我们比较了有无 LPC 预测器的 LRU 基线。结果表明,预测器的开销极小:当向 LRU 添加预测器时,SM Active 仅增加了 0.72%,SM Occupancy 增加了 0.87%,而 Tensor Pipe Active 仅增加了 0.17%。这证实了 LPC 预测器可以以微不足道的计算开销集成到 LPC 中。
A.2 LPC 内存开销敏感性分析
实验设置。我们通过将为预测器保留的内存按 0×、1×、2× 和 4× 进行缩放,来分析 LPC 性能对其 1 GB GPU 内存开销的敏感性。图 8 显示了在与第 4.3 节相同的设置下对命中率的影响。
结果与分析。结果显示,分配给 LPC 预测器的 GPU 内存对其命中率有 0.6% 到 8.1% 的影响。这是通过比较 LPC (0×) 和 LPC (1×) 观察到的。在缓存大小为 $16 \times 10^4$ tokens(40 GB for H100 after loading Qwen3-32B-FP8)时,命中率下降了 0.6% 到 2.4%。在较小的缓存容量下,命中率的影响更为明显,当从 LPC (0×) 变为 LPC (1×) 时,命中率下降了 2.3% 到 8.1%(在缓存大小为 $6.4 \times 10^4$ tokens,即 16 GB 时)。这些发现表明,尽管 LPC 的预测机制在根本上优于 LRU,但最小化其直接的 GPU 内存开销对于最大化其缓存命中率优势至关重要,尤其是在 GPU 内存资源有限的情况下。
LPC 与 LRU 对比。尽管存在内部开销变化,LPC 始终优于 LRU。当 LPC 预测器的开销从 0× 变化到 4× 时,在缓存大小为 $16 \times 10^4$ tokens (40 GB) 时,LPC 的命中率比 LRU 高 4.3% 到 16.7%。在较小的缓存大小 $6.4 \times 10^4$ tokens (16 GB) 时,LPC 的命中率比 LRU 高 5.8% 到 119.0%。这些发现表明 LPC 的预测机制在根本上优于 LRU。
图 8: LPC 缓存命中率对其内部 GPU 内存开销的敏感性,与 LRU 的比较。LPC (N×) 表示配置了 N 倍其实际 GPU 内存需求的 LPC。
A.3 LPC 预测器准确性与消融研究
评估指标。我们使用马修斯相关系数(MCC)和宏平均 F1 分数(F1 macro)来评估 LPC 预测器的有效性。MCC 是一个用于不平衡类别二元分类的稳健指标。F1 macro 为每个类别独立计算 F1 分数,然后取其平均值。LPCturns 是一项消融研究,仅使用轮次数作为输入,其中预测基于按轮次计数的训练数据分箱后的多数标签。
Table 3: LPC 预测器和一项仅使用轮次数作为输入的消融研究(LPCturns)的预测正确性。所有指标值越高越好。
结果分析。表 3 中的结果表明,完整的 LPC 预测器达到了 0.28 到 0.36 的 MCC,并且始终优于 LPCturns 消融研究。这表明利用对话的语义内容为预测继续可能性提供了有价值的信号。例如,在 Chatbot-Arena 数据集上,LPC 在 MCC(0.3689 vs. 0.0775)和 F1 macro(0.6837 vs. 0.4741)上比 LPCturns 有显著提升。ShareGPT 数据集上 LPCturns 的负 MCC 表明,对于该特定数据集,仅使用轮次计数作为预测器并不比随机猜测更好。LMSys 数据集在 LPC 和 LPCturns 之间的 MCC 和 F1 macro 相当,这表明在某些场景下,像 LPCturns 这样的轻量级预测器可能已经足够强大。总体而言,结果显示了我们内容感知特征在预测对话继续方面的有效性。
B 附加消融研究
B.1 输入特征消融:用户提示 vs. 模型响应
设计动机与实验。我们最初选择仅使用用户提示作为输入,是出于保持预测器轻量化的考虑。如表 4 所示,包含模型响应并未带来显著、一致的改进,甚至可能轻微降低性能。我们推测这是因为为追求效率而选择的紧凑型嵌入模型,可能缺乏从通常更长且更多样的模型响应中提取有用信号的能力,反而可能引入噪声。用户的提示似乎是其直接意图的更直接、更有力的信号。
Table 4: 预测器输入特征的消融研究:仅提示 vs. 提示加模型响应。性能通过马修斯相关系数(MCC)和 F1 宏观分数衡量。
B.2 提示窗口长度消融
设计选择与验证。我们选择了 5 个提示(当前提示 + 4 个先前提示)的默认历史窗口。这是基于我们的分析,即数据集中大多数对话都很短(例如,94-99% 的对话有五轮或更少)。为了验证这一选择,我们尝试将提示窗口长度增加到 10。表 5 显示这并未带来改善,并且通常导致性能略有下降。这再次证实,对于这些工作负载,最近的对话历史最具预测性。对于对话明显更长的工作负载,可以调整此超参数。
Table 5: 预测器提示窗口长度的消融研究。
B.3 泛化性:统一预测器 vs. 专用预测器
问题与实验。对于实际部署而言,一个关键问题是,单个通用预测器是否能与针对不同数据分布的专用预测器表现得一样好。我们通过在一个组合数据集上训练一个“统一”预测器,并将其性能与我们在每个测试集上的默认“专用”模型进行比较来测试这一点。表 6 的结果显示,统一模型在 Chatbot Arena 和 LMSys 数据集上具有竞争力。然而,其在 ShareGPT 上的性能明显较低,这表明虽然核心对话模式是可泛化的,但像 ShareGPT 这样的某些数据集可能包含独特的特性,从而受益于领域特定的训练。更大的模型可能会提高泛化性,但代价是更高的开销。
Table 6: 统一预测器与专用预测器的性能比较。
B.4 与 Oracle 的命中率比较
评估上限。为了评估进一步改进的潜力上限,我们将 LPC 的命中率与一个拥有完美未来知识(即知道一个对话是否会继续)的 Oracle 策略进行比较。图 9 展示了在三个数据集上,不同缓存大小下 LPC、LRU 和 Oracle 的比较。评估设置与第 4.3 节相同。
结果分析。结果表明,LPC 持续缩小了 LRU 和 Oracle 之间的命中率差距。在 LMSys 数据集(图 9a)上,差距缩小了 22% 到 30%。在 ShareGPT 数据集(图 9b)上,差距缩小了 28% 到 64%。在最小缓存尺寸下,与 Oracle 的剩余差距仅为当前 LPC 相对于 LRU 增益的约 0.5 倍。我们推测这是因为 LPC 在高置信度预测上准确性更高,而小缓存尺寸只够容纳那些高置信度的前缀。对于 Chatbot-Arena 数据集(图 9c),LPC 实现的从 LRU 到 Oracle 的命中率差距缩小范围为 25% 到 28%。
结论。这些观察结果凸显了 LPC 在前缀缓存方面取得的显著进步。虽然由 Oracle 基准定义的性能上限仍然存在,但我们认为 LPC 有效地捕捉了对话文本本身中绝大多数可用的预测信息。因此,我们相信未来要大幅缩小剩余差距并接近 Oracle 级别的性能,可能需要引入新的数据维度,如用户访问模式,而不仅仅是依赖文本内容进行预测。
图 9: 命中率 vs. 缓存大小。LPC 将 LRU 与 Oracle 之间的差距缩小了 22% 至 66%。
B.5 实验统计显著性
稳定性评估。为了评估我们结果的统计显著性和稳定性,我们为每个数据点重复实验运行至少 5 次。图 10 展示了带有误差条的命中率,误差条表示这些运行中观察到的最小值和最大值。对这些数据的分析显示,变化始终很小。具体来说,对于任何给定的配置(算法、数据集和缓存大小),观察到的最高和最低命中率之间的最大绝对差值小于 0.013。相对变化(最大-最小差值除以平均命中率)也很小,始终低于 4.6%,并且在大多数情况下低于 1%。这种在重复实验中低方差表明,我们提出的方法与基线之间观察到的性能差异是稳健的,并非由随机波动引起,从而强调了我们研究结果的统计显著性。
图 10: 带有误差条的命中率与缓存大小关系图。命中率的变化小于 0.013。
💬 评论讨论
欢迎在这里分享您的想法和见解!