论文

通过基于预测的 KV 缓存管理为动态代理工作流提供高效服务

Efficient Serving for Dynamic Agent Workflows with Prediction-based KV-Cache Management

模型推理KV Cache

摘要

基于 LLM 的工作流组成专门的代理来执行复杂的任务,并且这些代理通常共享大量上下文,允许 KV-Cache 重用以节省计算量。现有方法要么在代理级别管理 KV-Cache,但无法利用工作流中的重用机会,要么在工作流级别管理缓存,但假设每个工作流调用代理的静态序列。然而,实际的工作流程通常是动态的,其中调用代理的顺序以及由此引起的缓存重用机会取决于每个任务的上下文。为了有效地服务于这种动态工作流程,我们构建了一个名为 PBKV 的系统(\textbf{P}rediction-\textbf{B}ased \textbf{KV}-缓存管理)。对于每个工作流程,PBKV 通过融合历史工作流程和目标工作流程上下文的指导来预测未来几个步骤中的代理调用。根据预测,PBKV 估计缓存条目的重用潜力,并将高潜力条目保留在 GPU 内存中。为了对预测错误具有鲁棒性,PBKV 在缓存逐出和预取期间保守地利用预测。对三个工作流基准的实验表明,PBKV 在动态工作流上比 LRU 实现了高达 1.85\times$ 的加速,在静态工作流上比 SOTA 基线 KVFlow 实现了高达 $1.26\times$ 的加速。