Cache-Craft: Managing Chunk-Caches for Efficient Retrieval-Augmented Generation

📄 arXiv: 2502.15734v1 📥 PDF

作者: Shubham Agarwal, Sai Sundaresan, Subrata Mitra, Debabrata Mahapatra, Archit Gupta, Rounak Sharma, Nirmal Joshua Kapu, Tong Yu, Shiv Saini

分类: cs.DC, cs.AI, cs.CL, cs.LG, cs.OS

发布日期: 2025-02-05

备注: Accepted at SIGMOD 2025


💡 一句话要点

Cache-Craft:通过管理Chunk-Cache提升检索增强生成效率

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 检索增强生成 RAG 大型语言模型 LLM KV-cache 缓存管理 计算效率

📋 核心要点

  1. 现有RAG系统在处理重复检索的文本块时,LLM需要重复计算KV值,造成冗余计算和延迟增加。
  2. Cache-Craft通过管理和重用预计算的文本块KV值(chunk-caches),并结合少量重新计算来保证输出质量。
  3. 实验结果表明,Cache-Craft在减少冗余计算、提高吞吐量和降低延迟方面优于现有前缀缓存方法。

📝 摘要(中文)

检索增强生成(RAG)常与大型语言模型(LLM)结合使用,以注入领域知识或用户特定信息。在RAG中,给定用户查询,检索器从知识库中提取相关的文本块(chunks)。这些chunks作为输入提示的一部分发送给LLM。通常,任何给定的chunk会在不同的用户问题中被重复检索。然而,目前最先进的方法无法重用KV-cache,因为chunks出现在任意位置和任意上下文中,导致LLM的注意力层需要为输入chunks重复计算key values (KVs)。直接重用会导致输出质量下降。这导致了昂贵的GPU上潜在的冗余计算,并增加了延迟。本文提出了Cache-Craft,一个用于管理和重用RAG系统中与文本chunks对应的预计算KVs(称为chunk-caches)的系统。我们展示了如何识别可重用的chunk-caches,如何高效地执行少量重新计算以修复缓存以保持输出质量,以及如何在硬件中高效地存储和驱逐chunk-caches,以最大限度地提高重用率,同时掩盖任何开销。通过实际生产工作负载以及合成数据集,我们表明Cache-Craft比SOTA前缀缓存减少了51%的冗余计算,比完全重新计算减少了75%。此外,在实际生产工作负载上进行连续批处理时,对于LLaMA-3-8B和LLaMA-3-70B模型,我们在保持质量的同时,吞吐量提高了1.6倍,端到端响应延迟降低了2倍。

🔬 方法详解

问题定义:论文旨在解决RAG系统中由于重复检索相同文本块而导致的LLM计算冗余问题。现有方法,如前缀缓存,无法有效地重用这些文本块的KV-cache,因为它们在不同查询中的位置和上下文可能不同。这导致了不必要的计算开销和延迟增加。

核心思路:论文的核心思路是构建一个chunk-cache管理系统,该系统能够识别可重用的文本块,并利用预计算的KV值来避免重复计算。同时,为了保证输出质量,系统会执行少量的重新计算来适应不同的上下文。

技术框架:Cache-Craft系统包含以下主要模块:1) Chunk识别模块,用于识别输入中的文本块;2) Cache查找模块,用于在chunk-cache中查找匹配的文本块;3) Cache重用模块,用于重用预计算的KV值;4) 重新计算模块,用于执行少量重新计算以适应不同的上下文;5) Cache管理模块,用于存储和驱逐chunk-cache,以最大化重用率。

关键创新:Cache-Craft的关键创新在于其能够有效地管理和重用chunk-cache,即使文本块出现在不同的位置和上下文中。与现有方法相比,Cache-Craft通过少量重新计算来保证输出质量,避免了完全重新计算的开销。

关键设计:Cache-Craft的关键设计包括:1) 使用高效的索引结构来快速查找chunk-cache;2) 设计了一种重新计算策略,该策略只重新计算受上下文影响的部分KV值;3) 采用了一种基于LRU的cache驱逐策略,以最大化重用率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,Cache-Craft比SOTA前缀缓存减少了51%的冗余计算,比完全重新计算减少了75%。在实际生产工作负载上进行连续批处理时,对于LLaMA-3-8B和LLaMA-3-70B模型,Cache-Craft在保持质量的同时,吞吐量提高了1.6倍,端到端响应延迟降低了2倍。

🎯 应用场景

Cache-Craft可应用于各种基于RAG的LLM应用,例如问答系统、聊天机器人和文档摘要等。通过减少冗余计算和降低延迟,Cache-Craft可以提高这些应用的性能和用户体验,并降低部署成本。该研究对于提升LLM在实际应用中的效率具有重要价值。

📄 摘要(原文)

Retrieval-Augmented Generation (RAG) is often used with Large Language Models (LLMs) to infuse domain knowledge or user-specific information. In RAG, given a user query, a retriever extracts chunks of relevant text from a knowledge base. These chunks are sent to an LLM as part of the input prompt. Typically, any given chunk is repeatedly retrieved across user questions. However, currently, for every question, attention-layers in LLMs fully compute the key values (KVs) repeatedly for the input chunks, as state-of-the-art methods cannot reuse KV-caches when chunks appear at arbitrary locations with arbitrary contexts. Naive reuse leads to output quality degradation. This leads to potentially redundant computations on expensive GPUs and increases latency. In this work, we propose Cache-Craft, a system for managing and reusing precomputed KVs corresponding to the text chunks (we call chunk-caches) in RAG-based systems. We present how to identify chunk-caches that are reusable, how to efficiently perform a small fraction of recomputation to fix the cache to maintain output quality, and how to efficiently store and evict chunk-caches in the hardware for maximizing reuse while masking any overheads. With real production workloads as well as synthetic datasets, we show that Cache-Craft reduces redundant computation by 51% over SOTA prefix-caching and 75% over full recomputation. Additionally, with continuous batching on a real production workload, we get a 1.6X speed up in throughput and a 2X reduction in end-to-end response latency over prefix-caching while maintaining quality, for both the LLaMA-3-8B and LLaMA-3-70B models.