PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jiayi Yao, Hanchen Li, Yuhan Liu, Siddhant Ray, Yihua Cheng, Qizheng Zhang, Kuntai Du, Shan Lu, Junchen Jiang

arXiv:2405.16444v3cs.LG

TL;DR

Long RAG inputs reuse multiple text chunks, but non-prefix cached chunks omit cross-attention with preceding text, creating a speed–quality trade-off. CacheBlend selectively recomputes a small token subset while fusing cached chunks, reducing TTFT by 2.2–3.3× and increasing throughput by 2.8–5× without compromising generation quality. The evaluation supports faster multi-chunk KV-cache reuse, while the reported scope remains centered on transformer models and tested serving configurations.

  • Problem

    Existing KV-cache reuse either limits reuse to prefixes or ignores cross-attention among non-prefix chunks, making it difficult to combine multiple cached contexts with full-prefill quality.

  • Method

    CacheBlend fuses multiple precomputed KV caches by selectively recomputing a small fraction of token KV values using preceding texts from the specific input.

  • Results

    2.2–3.3× lower TTFT and 2.8–5× higher throughput than full KV recompute were achieved across four datasets and three models under negligible quality drop.

  • Takeaways & Limitations

    CacheBlend combines the speed benefits of KV-cache reuse with the generation quality of full KV recompute for inputs containing concatenated text chunks.

  • Takeaways & Limitations

    The method currently applies only to transformer models and was not evaluated across more models and datasets with different quantization settings.

Abstract

from arXiv · show

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.

1 Introduction

LLM applications often prepend multiple retrieved text chunks to queries, but long-input prefill makes generation slow and existing KV-cache reuse methods fail to combine non-prefix chunks without quality loss. CacheBlend selectively recomputes a small fraction of tokens to fuse multiple cached chunks, achieving faster inference while preserving generation quality.

  • RAG applications prepend multiple retrieved text chunks to user queries, adding domain or user-specific context for higher-quality responses.
  • Long inputs slow inference because prefill processes the entire input to produce the KV cache, and prefill delay determines TTFT.
  • Prefix caching reuses only the first input chunk, so inputs with many reused chunks remain nearly as slow as full KV recompute.
  • Full KV reuse handles non-prefix chunks by adjusting positional embeddings but ignores cross-attention among preceding chunks, which can be vital for jointly answering queries.
  • CacheBlend fuses multiple precomputed KV caches by selectively recomputing a small fraction of tokens based on preceding texts in the specific input.
  • CacheBlend’s partial KV update can be pipelined with cache fetching, enabling storage on slower devices without increasing inference latency.
  • 2.2–3.3× lower TTFT and 2.8–5× higher throughput than prefix caching were achieved without compromising generation quality or increasing storage cost.

2 Background

Transformer LLM serving first prefills input tokens into KV caches and then decodes iteratively, making long-context prefill a major latency and throughput concern. Prefix-cache reuse reduces computation by retaining cached prefix representations, but long inputs can still impose substantial delays.

  • The prefill phase transforms input tokens into key and value vectors that form the KV cache, after which decoding generates tokens iteratively.
  • Attention computation uses query and key vectors to form an attention matrix over preceding tokens, then combines it with value vectors to produce the layer output.
  • When a prefix KV cache is available, prefill computes only the forward attention between suffix tokens and prefix tokens on each layer.
  • A 4,000-token input can require three seconds for Llama-34B or six seconds for Llama-70B on one A40 GPU, creating substantial TTFT delay.
  • Prefill can also bottleneck throughput, with prior work showing that eliminating it can double LLM throughput.

3 Motivation

Repeated context chunks create substantial prefill overhead, motivating KV-cache reuse; however, prefix caching cannot reuse most chunks, while full KV reuse can miss cross-attention and reduce answer quality.

  • Motivation: Repeated context texts dominate prefill overhead, making KV-cache reuse attractive for recurring LLM inputs.The context portion is often larger than the user query, so recomputing it repeatedly contributes most of the prefill cost.
  • Why is prefix caching insufficient?: Prefix caching reuses only the first chunk when multiple chunks are prepended to an LLM input.Consequently, its savings are marginal when reused inputs contain several context chunks.
  • Why is prefix caching insufficient?: More retrieved chunks can improve generation quality, although too many chunks may hurt quality because of the lost-in-the-middle issue.The evaluation uses F1 score and retrieves top-k chunks by embedding distance.
  • Why is full KV reuse insufficient?: Full KV reuse preserves positional information for non-prefix chunks but omits their cross-attention with preceding chunks.The preceding text is unavailable when each chunk’s KV cache is precomputed.
  • Why is full KV reuse insufficient?: Missing cross-attention can change forward attention and produce wrong answers when a query requires jointly understanding multiple chunks.The paper illustrates this with Messi and Cristiano Ronaldo’s World Cup goal statistics.
  • Why is full KV reuse insufficient?: Full KV reuse can work when cross-attention between chunks is low, such as with prompt templates targeted by PromptCache.Its weakness is therefore most relevant to inputs whose chunks contain substantial interdependent information.

4 Fast KV Cache Fusing

CacheBlend targets fast, quality-preserving fusion of multiple precomputed KV caches by selectively recomputing only a small token subset on each layer. It chooses tokens using KV deviation patterns to reduce attention deviation toward full prefill.

  • Motivation: CacheBlend seeks both the speed of full KV reuse and the generation quality of full KV recompute for multi-chunk inputs.Its goal is to update precomputed KV caches so the resulting attention matrix and output text closely match full recompute.
  • Method: CacheBlend fuses multiple precomputed KV caches by recomputing a small fraction of tokens while reusing the others.The procedure runs layer by layer and updates only selected tokens based on the preceding texts in the specific input.
  • Terminology: KV deviation measures the absolute difference between a token-layer KV entry and the corresponding full-prefilled KV entry, while attention deviation measures the L-2 difference between attention matrices.These quantities guide token selection and assess how closely the fused cache approaches full recompute.
  • Selective KV recompute: Selective recomputation masks each layer’s input to selected tokens, computes their Q, K, and V vectors, then restores unselected KV entries before attention.The resulting attention still includes selected tokens attending to all tokens, while compute overhead scales with the selected-token count.
  • Token selection: Recomputing tokens with higher KV deviation reduces attention deviation more, with the largest reductions occurring among the highest-deviation tokens.These tokens are called High-KV-Deviation (HKVD) tokens.
  • Token selection: 10–20% HKVD-token recomputation suffices to greatly reduce attention deviation and preserve generation quality.The selection is practical because tokens with high KV deviation on one layer are likely to remain high-deviation on the next layer.

5 CacheBlend System Design

CacheBlend’s system design hides selective KV recomputation behind KV-cache loading, while controllers choose recomputation ratios and storage devices to avoid extra TTFT.

  • Pipelining KV loading and recompute: CacheBlend pipelines selective recomputation of one layer with loading the next layer’s precomputed KV cache into GPU memory.The loading delay can hide recomputation when loading is at least as slow as selective recomputation.
  • Pipelining KV loading and recompute: 3 ms per layer versus 16 ms for NVMe loading shows that recomputing 15% of tokens can incur no extra delay in the stated Llama-7B example.More tokens may also be recomputed without extra delay when recomputation remains below the loading delay.
  • Loading Controller: The loading controller selects a recomputation ratio by matching estimated recomputation and loading delays, subject to a minimum quality-preserving ratio.The minimum empirical ratio is 15%, even for faster storage devices.
  • Loading Controller: For a fixed 15% recomputation ratio, the controller chooses the cheapest storage device whose loading delay does not exceed recomputation delay.This allows storage cost to be reduced without adding inference delay.
  • KV cache store and Fusor: The cache store splits each LLM input into reusable or new text chunks, while the fusor merges their caches through layer-wise selective recomputation.The workflow connects retrieval, cache lookup, controller decisions, cache loading, and fusion before LLM inference.

6 Implementation

CacheBlend is implemented in vLLM through interfaces that fetch layer-specific caches, perform partial prefill, and synchronize layer processing.

  • Integration with vLLM: CacheBlend adds about 3K lines of Python code based on PyTorch v2.0 to vLLM.The integration exposes fetch_kv, prefill_layer, and synchronize interfaces for layer-wise partial prefill.
  • Integration with vLLM: prefill_layer performs the partial prefill for one layer and passes its output to the next layer.synchronize ensures that the current layer’s KV cache has reached GPU memory before prefilling.
  • Integration with vLLM: fetch_kv retrieves a text chunk’s KV cache by hashing the text and loading it from disk or CPU memory into GPU memory.Disk-backed caches use torch.load(), while CPU-resident caches use torch.cuda().

7 Evaluation

Across models, datasets, and RAG baselines, CacheBlend substantially reduces TTFT and raises throughput while keeping generation quality close to full KV recompute.

  • Overall improvement: 0.15–0.35 higher F1-score and Rouge-L scores than full KV reuse are reported, with no more than 0.01–0.03 quality drop versus full KV recompute and prefix caching.These results are summarized as evaluation takeaways across several models and tasks.
  • Evaluation setup: The evaluation spans Mistral-7B, Yi-34B, and Llama-70B, with datasets covering QA, summarization, and synthetic RAG chunk reuse.Reported datasets include 2WikiMQA, Musique, SAMSum, MultiNews, and extended RAG variants.
  • Overall improvement: 2.2–3.3× lower TTFT is achieved versus full KV recompute across several models and tasks, with quality reduction within 0.02 against full recompute and prefix caching.CacheBlend’s quality also exceeds full KV reuse by a large margin in many cases.
  • Overall improvement: 2.8–5× higher throughput is achieved than the baselines across different models and datasets while maintaining lower delay.The comparison covers full KV recompute and prefix caching under different request rates.
  • Comparison with RAG methods: CacheBlend has 2–5× lower TTFT and higher F1 score than MapReduce, while MapRerank has slightly lower TTFT but much worse quality.The paper attributes MapReduce’s delay to additional LLM inference and MapRerank’s quality loss to separate chunk processing.
  • Configuration analysis: CacheBlend maintains minimal quality degradation when KV caches are stored in either RAM or a slower SSD device.The storage-device experiment uses Yi-34B on 2WikiMQA.

8 Related Work

CacheBlend complements prior work on context reuse, general-purpose serving, and context compression by addressing multi-chunk KV-cache fusion.

  • RAG: RAG uses externally fetched text chunks to improve LLM accuracy and reliability, but processing those chunks can increase inference time.CacheBlend stores and reuses their KV caches to reduce this overhead.
  • KV cache reuse across requests: Prior KV-cache reuse work mostly focuses on prefixes, while PromptCache supports different positions but loses quality through inaccurate positional encoding and ignored cross-attention.CacheBlend uses partial recomputation to retain positional accuracy and cross-attention information.
  • General-purpose LLM serving systems: CacheBlend is complementary to general-purpose serving systems such as Orca and vLLM, adding context-reuse capabilities to their serving functions.These systems address parallel processing and GPU memory management rather than the same cache-fusion problem.
  • Context compression methods: CacheBlend is compatible with context-shortening methods and can benefit from techniques that drop unimportant KV vectors by storing and loading less cache.The paper also reports compatibility with different chunk lengths.

9 Limitations and Future Work

The paper identifies scope and integration boundaries for CacheBlend. Its current method is limited to transformer models, selected evaluation settings, and the vLLM serving engine.

  • Scope: CacheBlend currently applies only to language models with transformer structures.Architectures such as Mamba and Griffin are left for future investigation.
  • Evaluation boundaries: The evaluation does not cover more models and datasets under different quantization settings.
  • Integration: CacheBlend has been integrated with vLLM but not tested on newer serving engines such as Distserve or StableGen.
  • Integration: Applying CacheBlend to workloads sharing KV caches across different compute nodes remains unstudied.Integration with these inference frameworks is left for future work.

10 Conclusion

CacheBlend combines precomputed KV caches from concatenated text chunks and selectively recomputes a small fraction of tokens to recover cross-attention. Across four datasets and three models, it reduces latency and increases throughput with negligible quality drop.

  • CacheBlend combines multiple precomputed KV caches when their corresponding texts are concatenated in the LLM input.
  • CacheBlend selectively recomputes KV values for a small fraction of tokens to preserve generation quality by recovering cross-attention among texts.
  • 2.2-3.3× lower TTFT and 2.8-5× higher throughput than full KV recompute were observed across four datasets and three models.The reported comparison comes with negligible quality drop.

A N-dimensional positional recovery

This section defines RoPE in N-dimensional space and proves that its attention score depends on relative rather than absolute token positions. That property supports positional recovery for reused KV caches.

  • The section begins by defining RoPE for query and key vectors embedded at position m in N-dimensional space.
  • RoPE encodes positional information through a rotary matrix whose hyperparameters are given by Θ.
  • The positional recovery rationale is that attention scores between token pairs are invariant to their absolute positions.
  • The proof derives the attention score between a key at position m and a query at position m+l using the RoPE transformations.
  • The resulting attention score depends only on relative distance l, not absolute position m.
Loading 2405.16444v3…