PaperlayerView exact revision on arXiv

Source-linked AI summary

TurboRAG: Accelerating Retrieval-Augmented Generation with Precomputed KV Caches for Chunked Text

Songshuo Lu, Hua Wang, Yutian Rong, Zhi Chen, Yaohua Tang

arXiv:2410.07590v1cs.CVcs.CL

TL;DR

Conventional RAG’s concatenated retrieved chunks make prefill computationally costly and increase time-to-first-token. TurboRAG precomputes document KV caches for online reuse and adapts attention and positional handling; across LongBench multi-document QA benchmarks, it lowers TTFT while retaining comparable accuracy.

  • Problem

    Conventional RAG recomputes KV caches for retrieved chunks and processes long concatenated inputs, increasing prefill cost and TTFT.

    Source: §1 · p. 1

  • Method

    TurboRAG precomputes document KV caches offline and reuses them during prefill with independent attention and positional handling.

    Source: §1 · p. 2, §3.2 · p. 5

  • Results

    Across LongBench multi-document QA benchmarks, TurboRAG reduces TTFT by up to 9.4x and on average 8.6x, with comparable accuracy to the conventional RAG baseline.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Online inference uses 98.46% fewer computational resources than standard RAG, increasing maximum supported batch size and throughput.

    Source: §1 · p. 2

    The new attention mask and position embedding can significantly reduce question-answering accuracy, so the model needs specific training for this setting.

    Source: §3.2 · p. 5

Abstract

from arXiv · show

Current Retrieval-Augmented Generation (RAG) systems concatenate and process numerous retrieved document chunks for prefill which requires a large volume of computation, therefore leading to significant latency in time-to-first-token (TTFT). To reduce the computation overhead as well as TTFT, we introduce TurboRAG, a novel RAG system that redesigns the inference paradigm of the current RAG system by first pre-computing and storing the key-value (KV) caches of documents offline, and then directly retrieving the saved KV cache for prefill. Hence, online computation of KV caches is eliminated during inference. In addition, we provide a number of insights into the mask matrix and positional embedding mechanisms, plus fine-tune a pretrained language model to maintain model accuracy of TurboRAG. Our approach is applicable to most existing large language models and their applications without any requirement in modification of models and inference systems. Experimental results across a suite of RAG benchmarks demonstrate that TurboRAG reduces TTFT by up to 9.4x compared to the conventional RAG systems (on an average of 8.6x), but reserving comparable performance to the standard RAG systems.

1 INTRODUCTION

TurboRAG addresses the high prefill cost and latency of conventional RAG by reusing document KV caches computed offline. Across LongBench multi-document QA benchmarks, it reduces TTFT while maintaining comparable accuracy to the baseline.

1 INTRODUCTION

  • Conventional RAG recomputes KV caches for recalled chunks and incurs quadratic prefill cost as augmented sequences grow, increasing TTFT and potentially limiting batch size.TurboRAG’s hybrid offline-online prefill is proposed to eliminate much of the online computation.Source: §1 · pp. 1–2
  • TurboRAG precomputes and stores each document’s KV cache offline, then injects retrieved caches into requests using independent attention and standard RoPE.The approach is motivated by sparse cross-attention between documents and the equivalence of relative positional embeddings for individual versus concatenated documents.Source: §1 · p. 2
  • Up to 9.4x lower TTFT across LongBench multi-document QA benchmarks, with comparable accuracy to the conventional RAG baseline.The reported average TTFT reduction is 8.6x.Source: §1 · p. 2
  • The paper’s contributions combine offline-online prefill with attention-mask and position-ID techniques intended to maintain accuracy while reducing KV-cache computation overhead.The reported headline improvement is 9.4x in TTFT over multi-document QA benchmarks without compromising accuracy.Source: §1 · p. 2

2 RELATED WORK

Prior RAG work covers retrieval-and-generation workflows, attention and KV-cache optimization, and distributed inference. TurboRAG targets a gap specific to dynamic multi-document retrieval: offline reuse of document-fragment caches while addressing accuracy and concatenation challenges.

2 RELATED WORK

  • Standard RAG embeds and indexes document chunks offline, then retrieves relevant chunks and concatenates them with the user request for generation.The workflow combines offline preparation with online retrieval and generation.Source: §2 · p. 2
  • Prior optimization work modifies attention computation, compresses or merges KV caches, or targets distributed large-model inference.The cited approaches address general long-text generation or distributed inference rather than TurboRAG’s RAG-specific setting.Source: §2 · p. 3
  • Dynamic retrieved fragments make direct KV-cache concatenation an accuracy concern, while a prior multilevel cache reduces redundant computation without analyzing model accuracy.This work is identified as Jin et al. (2024a).Source: §2 · p. 3
  • TurboRAG precomputes and stores document-fragment KV caches offline for direct online use, redesigning RAG prefill to avoid recomputing retrieved fragments.The paper presents this as a way to reduce computational complexity for latency-constrained applications.Source: §2 · p. 3

3 METHODOLOGY

TurboRAG formalizes prefill around document-level KV caches, adjusts attention masks and position IDs, and trains the model for those inference patterns. Its method specifies how to form and use the caches across offline and online stages.

3 METHODOLOGY

  • The methodology introduces TurboRAG and sets out its formalization, attention and position changes, model adaptation, and inference pipeline.Figure 2 contrasts causal attention with independent attention using composite or reordered positions.Source: §3 · pp. 3–4, Figure 2 · p. 4

3.1 PROBLEM FORMALIZATION

  • In conventional RAG, retrieved chunks are concatenated with the query for prefill, producing context hidden states used to compute query, key, and value matrices.The resulting KV cache is stored for subsequent attention computation during autoregressive decoding.Source: §3.1 · pp. 4–5

3.2 POSITION ID REARRANGEMENT

  • Independent Attention masks cross-document attention, while reordered positions restore monotone token positions across concatenated chunks.The query, answer, and all documents still participate in attention during decoding.Source: Figure 2 · p. 4, §3.2 · p. 5

3.3 ADAPTING LLMS FOR PRECOMPUTED CACHE CONCATENATION

  • TurboRAG fine-tunes the LM with the attention masks and position embeddings used at inference to avoid question-answering accuracy regression.The method contrasts these patterns with standard supervised fine-tuning using the Figure 2a setting.Source: §3.2 · p. 5, §3.3 · p. 6

3.4 THE TURBORAG PIPELINE

  • Offline, the LM computes and stores each document's KV cache; online, TurboRAG retrieves and concatenates caches, prefills the query, then generates a response.The pipeline also encodes documents and queries for vector-search retrieval.Source: §3.4 · p. 6

4 EXPERIMENTS

The experiments evaluate document QA accuracy, general-capability regression, TTFT, and batch scaling across TurboRAG variants and conventional RAG. TurboRAG-reordered combines comparable accuracy with substantial TTFT gains.

4.1 EXPERIMENT SETUP

  • The study compares Naïve RAG with composite-position and reordered-position TurboRAG variants, fine-tuned on data comprising 50% document QA and 50% general tasks.The training base was Qwen2-7B.Source: §4.1 · p. 7

4.2 DOCUMENT QA ACCURACY

  • After fine-tuning, TurboRAG-reordered and TurboRAG-composite kept the RGB accuracy gap within 1% of Naïve RAG and performed comparably to GPT-4o across Chinese and English, including high-noise conditions.Without fine-tuning, accuracy dropped significantly as noise increased.Source: §4.2 · p. 7, Table 1 · p. 8
  • TurboRAG achieved answer accuracy comparable to Naïve RAG on LongBench multi-document QA subcategories.Source: §4.2 · p. 7, Table 2 · p. 8

4.3 GENERAL CAPABILITY REGRESSION

  • OpenCompass regression tests showed minimal impact on base capabilities, with TurboRAG-reordered exhibiting no significant performance degradation versus Naïve RAG.Source: Table 3 · p. 8, §4.3 · p. 8

4.4 TTFT PERFORMANCE

  • 8.6x average TTFT speedup, peaking at 9.4x, was reported for TurboRAG-reordered versus Naïve RAG on LongBench multi-document QA.The experiments used Hugging Face Transformers, FlashAttention2, and an NVIDIA A100 80GB GPU.Source: §4.4 · p. 9

4.5 BATCH SCALING

  • Table 4 shows TTFT speedup increasing with batch size; with host-to-device KV-cache transfer, TurboRAG remained fourfold faster and used approximately 98.46% fewer TFLOPs than Naïve RAG.Source: Table 4 · p. 9, §4.5 · p. 9

5 CONCLUSION AND DISCUSSION

TurboRAG reduces prefill time by changing how retrieved text fragments are processed. The discussion also highlights model adaptability to different positional embeddings.

5 CONCLUSION AND DISCUSSION

  • TurboRAG reduces the time required for prefill computations when concatenating retrieved text fragments.The authors describe KV-cache compression as orthogonal and directly usable alongside their approach to reduce latency and ease storage pressure.Source: §5 · p. 10
  • The authors raise whether cross-attention between fragments is necessary and suggest that some cross-attention computation for long-text inference might be redundant.Their analogy is that separate people can each provide information to a questioner without communicating with one another.Source: §5 · p. 10
  • The authors report that the model adapted to two different sets of positional embeddings in their work.They relate this to context-window extension by position interpolation, where pretrained models encounter a longer context length during continued training.Source: §5 · p. 10

A DOCUMENT Q&A EXAMPLE

The example asks when Carole King & James Taylor: Just Call Out My Name premiered; the retrieved documents include the answer alongside other topics.

A DOCUMENT Q&A EXAMPLE

B DATA PROPORTIONS

The appendix includes tables titled for model fine-tuning data sampling ratios and document-Q&A data and quantities.

B DATA PROPORTIONS

  • Table 5 concerns sampling ratios for different data types used during model fine-tuning.Source: Table 5 · p. 14
  • Table 6 concerns the specific data and quantities used for document question answering.Source: Table 6 · p. 14

C COMPUTATIONAL LOAD CALCULATION

The section describes a FLOPS calculation for a Swiglu LLM, summing per-layer costs for QKV transformation, attention, projection, and the MLP across input tokens and layers.

C COMPUTATIONAL LOAD CALCULATION

Related papers

Loading 2410.07590v1…