CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion

TL;DR

CacheBlend tackles the slow prefill bottleneck in LLMs when inputs include multiple retrieved text chunks by fusing independently precomputed KV caches with selective KV recomputation. By updating only a small fraction of tokens per layer (HKVD-based selection) and pipelining recompute with KV loading, it preserves cross-attention and generation quality while achieving large reductions in time-to-first-token and increases in throughput. Evaluations across three open-source LLMs and four tasks show TTFT reductions of 2.2-3.3x and throughput gains of 2.8-5x with negligible quality loss compared to full recompute, and significantly better latency/quality trade-offs than prefix caching or full KV reuse. The approach is implemented on top of vLLM and is designed to work with multi-chunk RAG contexts, enabling scalable, storage-friendly KV-cache serving for long-context inference.

Abstract

Large language models (LLMs) often incorporate multiple text chunks in their inputs to provide the necessary contexts. To speed up the prefill of the long LLM inputs, one can pre-compute the KV cache of a text and re-use the KV cache when the context is reused as the prefix of another LLM input. However, the reused text chunks are not always the input prefix, which makes precomputed KV caches not directly usable since they ignore the text's cross-attention with the preceding texts. Thus, the benefits of reusing KV caches remain largely unrealized. This paper tackles just one challenge: when an LLM input contains multiple text chunks, how to quickly combine their precomputed KV caches in order to achieve the same generation quality as the expensive full prefill (i.e., without reusing KV cache)? This challenge naturally arises in retrieval-augmented generation (RAG) where the input is supplemented with multiple retrieved texts as the context. We present CacheBlend, a scheme that reuses the precomputed KV caches, regardless prefix or not, and selectively recomputes the KV values of a small subset of tokens to partially update each reused KV cache. In the meantime, the small extra delay for recomputing some tokens can be pipelined with the retrieval of KV caches within the same job, allowing CacheBlend to store KV caches in slower devices with more storage capacity while retrieving them without increasing the inference delay. By comparing CacheBlend with the state-of-the-art KV cache reusing schemes on three open-source LLMs of various sizes and four popular benchmark datasets of different tasks, we show that CacheBlend reduces time-to-first-token (TTFT) by 2.2-3.3x and increases the inference throughput by 2.8-5x from full KV recompute without compromising generation quality. The code is available at https://github.com/LMCache/LMCache.

Figures (17)

• Figure 1: Contrasting full KV recompute, prefix caching, full KV reuse, and CacheBlend's selective KV recompute.
• Figure 2: Generation quality improves as more text chunks are retrieved.
• Figure 3: An illustrative example of an LLM input with two text chunks prepended to a query. Full KV recompute (b), without reusing KV cache, is slow but gives the correct answer. Full KV reuse (c), however, gives the wrong answer as it neglects cross-attention between the chunks (Figure \ref{['fig:attn_comparison']}).
• Figure 4: Contrasting the attention matrices of (a) full KV recompute and (b) full KV reuse. The yellow boxes highlight the cross-attention. The right-hand side plots show the resulting forward attention matrices whose discrepancies are a result of the different cross-attention between the two methods.
• Figure 5: Illustrated contrast between (a) full KV recompute and (b) selective KV recompute on one layer.

Theorems & Definitions (2)

• Definition 1: Rotary Positional Encoding, ROPEsu2024roformer
• proposition 1: Rope only depends on relative position