Source-linked AI summary

RAGCache: Efficient Knowledge Caching for Retrieval-Augmented Generation

Chao Jin, Zili Zhang, Xuanlin Jiang, Fangyue Liu, Xin Liu, Xuanzhe Liu, Xin Jin

arXiv:2404.12457v2cs.DCcs.CLcs.LG

TL;DR

RAG’s injected knowledge creates long sequences that increase computation and memory costs. RAGCache addresses this with multilevel knowledge-state caching, retrieval-aware replacement, and speculative overlap of retrieval and inference, achieving up to 4× lower TTFT and 2.1× higher throughput than vLLM with Faiss.

  • Problem

    Knowledge injection lengthens RAG sequences, creating high computation and memory costs and motivating more efficient processing.

  • Method

    RAGCache caches retrieved-document intermediate states in a GPU–host knowledge hierarchy, uses retrieval-aware replacement, and overlaps retrieval with inference.

  • Results

    Up to 4× lower TTFT and 2.1× higher throughput are reported versus vLLM integrated with Faiss.

  • Takeaways & Limitations

    Caching repeated knowledge states and overlapping retrieval with inference reduce redundant computation and end-to-end latency in evaluated RAG workloads.

  • Takeaways & Limitations

    The speculative-pipelining analysis assumes one request at a time for both vector search and the LLM, with batch size one.

Abstract

from arXiv · show

Retrieval-Augmented Generation (RAG) has shown significant improvements in various natural language processing tasks by integrating the strengths of large language models (LLMs) and external knowledge databases. However, RAG introduces long sequence generation and leads to high computation and memory costs. We propose RAGCache, a novel multilevel dynamic caching system tailored for RAG. Our analysis benchmarks current RAG systems, pinpointing the performance bottleneck (i.e., long sequence due to knowledge injection) and optimization opportunities (i.e., caching knowledge's intermediate states). Based on these insights, we design RAGCache, which organizes the intermediate states of retrieved knowledge in a knowledge tree and caches them in the GPU and host memory hierarchy. RAGCache proposes a replacement policy that is aware of LLM inference characteristics and RAG retrieval patterns. It also dynamically overlaps the retrieval and inference steps to minimize the end-to-end latency. We implement RAGCache and evaluate it on vLLM, a state-of-the-art LLM inference system and Faiss, a state-of-the-art vector database. The experimental results show that RAGCache reduces the time to first token (TTFT) by up to 4x and improves the throughput by up to 2.1x compared to vLLM integrated with Faiss.

1 Introduction

RAG improves generation by injecting retrieved external knowledge, but this creates long augmented sequences and high computation and memory costs. RAGCache addresses these costs by caching reusable knowledge states and overlapping retrieval with inference.

  • 1 Introduction: Retrieved documents improve generation quality, but their injection substantially lengthens requests and raises computation and memory costs.A 100-token request with 1000 retrieved tokens incurs more than 10× the original computation and memory costs.
  • 1 Introduction: RAGCache caches intermediate states of frequently retrieved documents and shares them across requests to reduce redundant computation.Its system characterization identifies repeated documents and skewed retrieval frequency as the caching opportunity.
  • 1 Introduction: A knowledge tree places frequently accessed states in GPU memory and less frequently accessed states in host memory, while retrieval order affects reusable states.Different document orders produce different key-value tensors because each new token depends on preceding tokens.
  • 1 Introduction: Dynamic speculative pipelining overlaps CPU retrieval with GPU inference because retrieval and inference otherwise execute sequentially.The strategy aims to reduce end-to-end latency while keeping system load under control.
  • 1 Introduction: Up to 4× lower TTFT and 2.1× higher throughput versus vLLM with Faiss establish RAGCache’s main benchmark result.It also achieves up to 3.5× lower TTFT and 1.8× higher throughput than SGLang.

2 Background

RAG combines offline vector indexing and online retrieval with GPU-based generation. Its end-to-end performance depends on retrieval cost, model size, and the length of the augmented sequence.

  • 2 Background: RAG improves responses by dynamically retrieving external information and combining it with the generative model’s capabilities.The workflow is intended to produce more accurate, relevant, and contextually rich responses.
  • 2 Background: RAG transforms documents into vectors offline, indexes them, and retrieves the most relevant documents through online vector similarity search.The retrieved documents are then used in the generation workflow.
  • 2 Background: RAG uses CPU-based retrieval and GPU-based generation, so end-to-end performance depends on both retrieval and generation.Retrieval time is mainly determined by vector-database scale, while generation time depends on model size and sequence length.

3 RAG System Characterization

Characterization identifies LLM generation, especially long-sequence prefill, as RAG’s main bottleneck and shows that caching retrieved knowledge can reduce redundant computation. Retrieval patterns and memory-transfer costs determine when caching is effective.

  • 3 RAG System Characterization: The characterization evaluates bottlenecks, caching benefits, and retrieval patterns to identify optimization opportunities for RAG.It focuses on LLM generation, intermediate-state caching, and question-pattern analysis.
  • 3 RAG System Characterization: An average document length of 3718 tokens makes generation markedly slower than retrieval in most cases, although high-accuracy search can make retrieval comparable.The generation sequence includes both original-request and retrieved-document tokens.
  • 3 RAG System Characterization: Caching retrieved documents’ key-value tensors reduces prefill latency because subsequent requests can reuse intermediate states instead of recomputing them.With a cached prefix, only the request tokens’ key-value tensors need computation.
  • 3 RAG System Characterization: Cache effectiveness depends on hit rate and memory hierarchy: GPU capacity is limited, while host-memory storage adds key-value-cache transfer overhead.Caching is ineffective when every request retrieves a unique document.
  • 3 RAG System Characterization: Retrieval patterns remain similar across additional embedding models and approximate-nearest-neighbor indexes, supporting caching optimization across settings.The results are consistent with the FlatL2 index.

4 RAGCache Overview

RAGCache reduces redundant computation by caching retrieved documents’ key-value tensors across requests. Its knowledge tree, PGDSF policy, and RAG controller coordinate cache organization, eviction, retrieval, and generation.

  • RAGCache caches retrieved documents’ key-value tensors across multiple requests to minimize redundant computation.
  • The RAG controller coordinates database retrieval with cache lookup and LLM generation, using cache-aware reordering and dynamic speculative pipelining.Dynamic pipelining can initiate LLM inference from retrieval results produced mid-process.
  • A knowledge tree organizes cached tensors as document-prefix paths, preserving document order and enabling shared prefixes across request sequences.
  • PGDSF retains valuable tensors by considering access frequency, tensor size, last access time, and prefix-aware recomputation cost.

5 RAGCache Design

RAGCache organizes order-sensitive document KV states in a multilevel knowledge-tree cache and manages them with prefix-aware replacement, reordering, and retrieval–inference overlap.

  • 5 RAGCache Design: RAGCache combines a knowledge tree, prefix-aware PGDSF replacement, cache-aware request reordering, and dynamic speculative pipelining for RAG.These components target document-order sensitivity, cache thrashing, cache placement, and retrieval latency.
  • 5.1 Cache Structure and Replacement Policy: Document order changes KV tensors because each token’s representation depends on preceding tokens, so RAGCache caches ordered document prefixes rather than individual documents.The knowledge tree maintains document order while enabling efficient retrieval of cached states.
  • 5.1 Cache Structure and Replacement Policy: PGDSF prioritizes nodes using recency, retrieval frequency, tensor size, and prefix-aware recomputation cost, evicting lower-priority nodes across GPU and host memory.Its cost estimation uses profiled prefill times and bilinear interpolation for cached and non-cached token lengths.
  • 5.1 Cache Structure and Replacement Policy: RAGCache swaps a node’s KV tensors to host memory only on first eviction, then frees later GPU copies without repeated transfers.This reduces PCIe data movement while using host memory’s substantially larger capacity.
  • 5.2 Cache-aware Reordering: Cache-aware reordering prioritizes requests with larger cached portions relative to computation needs, increasing cache hits while using fairness windows to prevent starvation.The policy addresses alternating requests that otherwise repeatedly swap cached documents and can produce zero cache hits.
  • 5.3 Dynamic Speculative Pipelining: Dynamic speculative pipelining overlaps vector retrieval with LLM inference by using early retrieval results for speculative generation and replacing speculation when later results differ.If the final retrieved documents match the latest speculation, the generated result can be returned directly.

6 Implementation

RAGCache is implemented as a vLLM-based prototype with prefix-caching support and Faiss-based dynamic speculative pipelining for IVF and HNSW indexes.

  • 6 Implementation: The prototype contains ∼5000 lines of C++ and Python and extends vLLM v0.3.0 for prefix caching across multiple attention mechanisms.The implementation extends prefill kernels in PyTorch and Triton for multi-head and grouped-query attention.
  • 6 Implementation: Dynamic speculative pipelining is implemented on Faiss and adapted to IVF and HNSW vector indexes.IVF searches within nearby clusters, while HNSW represents the vector space with multi-level graphs.
  • 6 Implementation: RAGCache replicates frequently accessed upper-level knowledge-tree nodes in host memory to recover from GPU failures and retries failed request processing with timeouts.Requests failing before their first iteration are recomputed; later failures can continue computation.

7 Evaluation

RAGCache is evaluated across models, datasets, retrieval settings, cache policies, reordering, and speculative pipelining, consistently reducing latency and improving throughput over vLLM and SGLang.

  • 7.1 Overall Performance: 1.2–4× lower average TTFT than vLLM and 1.1–3.5× lower than SGLang yields 1.3–2.1× and 1.2–1.8× higher throughput, respectively.The advantage holds across Mistral-7B and LLaMA2-7B on MMLU and Natural Questions.
  • 7.2 Case Study: RAGCache outperforms vLLM by 1.7–3.1× and SGLang by 1.2–2.5× in average TTFT across top-k values of 1, 3, and 5.The knowledge tree evicts the node furthest from the root, preserving frequently used prefixes despite factorially growing document permutations.
  • 7.2 Case Study: RAGCache maintains TTFT below 1.4 seconds across request rates for Mixtral-8×7B and LLaMA2-70B, while vLLM misses its SLO above 2 and 1.5 requests per second.RAGCache also surpasses SGLang by 1.2–2.6× in average TTFT on the large-model workloads.
  • 7.3 Ablation Study: PGDSF achieves 1.02–1.32× higher hit rate than GDSF, 1.06–1.62× than LRU, and 1.06–1.75× than LFU, lowering average TTFT by 1.05–1.29×.The policy accounts for document-prefix sizes, access patterns, and recomputation costs.
  • 7.3 Ablation Study: Cache-aware reordering reduces average TTFT by 1.2–2.1× under high request rates, while dynamic speculative pipelining reduces TTFT by up to 1.6×.Speculative pipelining also decreases non-overlapping vector-search time by 1.5–4.3×.

8 Discussion

RAGCache targets RAG’s long prefills by caching document KV states and reusing shorter prefixes when retrieved-document permutations become difficult to reuse.

  • Time per output token (TPOT): Caching frequently retrieved document KV states primarily reduces RAG’s prolonged TTFT caused by knowledge-increased input length.RAGCache can also lower TPOT because caching accelerates the prefill iteration, which typically dominates decoding iterations.
  • Large top-k: As top-k increases, factorially many document permutations reduce reuse, so RAGCache caches lower-top-k prefixes such as top-3 prefixes for top-5 requests.This balances cache hit rate against cache efficiency.

9 Related Work

RAGCache extends prior RAG, vector-search, KV-cache, and KV-cache-reuse techniques with retrieval-pattern-aware caching and retrieval–inference pipelining.

  • RAG: RAGCache supports iterative retrieval by treating intermediate iterations as separate requests and caching their documents’ corresponding KV states.
  • Vector search: RAGCache extracts temporary vector-search results for speculative LLM generation, thereby pipelining vector search with LLM inference.
  • KV cache management: Unlike KV-cache quantization, compression, and token-subset methods that introduce approximation, RAGCache preserves exact document KV states without affecting generation quality.
  • KV cache reusing: RAGCache builds a multilevel cache around RAG retrieval patterns, whereas prior reuse systems may reuse tokens flexibly, compress KV states, or reuse GPU-resident caches.

10 Conclusion

RAGCache combines multilevel caching, prefix-aware replacement, and dynamic speculative pipelining to reduce redundant RAG computation and overlap retrieval with inference.

  • 10 Conclusion: RAGCache outperforms vLLM integrated with Faiss by up to 4× on TTFT and 2.1× on throughput across varied models and workloads.
Loading 2404.12457v2…