SIFT: Selective-Index For Fast Compute of RAG Prefill by Exploiting Attention Invariance

2026-06-08Artificial Intelligence

Artificial IntelligenceHardware Architecture
AI summary

The authors study how to speed up a technique called Retrieval-Augmented Generation (RAG), which improves large language model answers by adding relevant documents but slows down response time. They find that since many RAG queries reuse the same documents, it’s wasteful to fully recompute everything each time. Their method, SIFT, identifies small important parts of the documents based on attention patterns and only recomputes those, drastically reducing data storage and speeding up answers without much loss in accuracy. This works by leveraging stable attention spots inside and across documents, allowing fast and efficient reuse.

Retrieval-Augmented GenerationLarge Language ModelsPrompt LengthTime to First TokenKey-Value TensorsAttention InvarianceLocal-AttentionCross-AttentionPrefillDisk Transfer Latency
Authors
Rya Sanovar, Srikant Bharadwaj, Hritvik Taneja, Moinuddin Qureshi
Abstract
Retrieval-Augmented Generation (RAG) injects LLM queries with relevant documents to improve response quality. This injection increases prompt length and slows time to first token (TTFT). Unlike standard queries, RAG queries have a unique property of context reuse where the same documents recur across user queries. Thus, fully recomputing documents for every RAG query does redundant compute and increases TTFT. Prior works precompute KV tensors of RAG documents offline and coarsely recompute some tokens during online prefill. However, such KV reuse is often slower than full recomputation on modern GPUs due to high-latency disk transfers. Further, such a coarse-grained recomputation degrades accuracy. To address these limitations, this paper proposes SIFT: Selective-Index For Fast Compute of RAG Prefill by Exploiting Attention Invariance. SIFT processes documents offline and extracts fine-grained locations of high attention scores for each document. Next, we identify the following attention invariance insights that enable us to exploit the extracted locations during runtime: (1) Local-Attention Invariance: The location of high attention scores within a document remain invariant to surrounding documents. This helps us predict the location of high scores where the document attends to itself. (2) Cross-Attention Consistency: Keys with high intra-document attention also attract cross-attention from subsequent documents. This helps us predict the location of high scores where the document attends to future documents. Critically, SIFT stores no KV data and only stores locations of high scores in the form of two compact bit vectors. SIFT's storage is up to 24,000x smaller than KV tensors, obviating costly disk transfers. During prefill, SIFT computes the attention only for the marked locations and improves TTFT by 1.71x while holding accuracy within 1% of full recompute.