LazyAttention: Efficient Retrieval-Augmented Generation with Deferred Positional Encoding
Abstract
Key-value (KV) caching accelerates inference of large language models (LLMs) by reusing past computations for generated tokens. Its importance becomes even greater in long-context applications such as retrieval-augmented generation (RAG) and in-context learning (ICL). However, conventional KV caching embeds positional information directly into the cache, limiting its reusability. Existing solutions either restrict reuse to prefixes or require expensive memory materialization for positional re-encoding. We introduce LazyAttention, a novel attention mechanism that kernelizes deferred positional encoding to enable zero-copy, position-agnostic KV reuse. By adjusting positional encoding within attention kernels on-the-fly, LazyAttention resolves the materialization bottleneck, allowing a single physical KV copy to serve multiple logical requests at arbitrary positions. Leveraging attention kernels tailored for prefilling and decoding, our system achieves significant efficiency improvements: under skewed document distributions, it reduces time-to-first-token (TTFT) by 1.37× and increases inference throughput by 1.40× compared to the state-of-the-art Block-Attention, while maintaining comparable output quality.
Lay Summary
Large language models are increasingly used with retrieval, where they answer questions by first reading relevant documents. This improves accuracy and freshness, but it also creates a speed problem: the model often has to reprocess the same documents again and again, especially when they appear in different positions in a prompt. We built LazyAttention, a new attention mechanism that lets the model reuse stored document computations no matter where those documents appear. Instead of permanently baking position information into the cache, LazyAttention adds position information only at the moment attention is computed. This avoids copying or rebuilding the cache while keeping the model’s calculations effectively the same. This makes retrieval-augmented generation faster and more memory-efficient. In our experiments, LazyAttention reduced time-to-first-token by up to 1.37× and increased throughput by 1.40× compared with a strong recent baseline, while preserving similar answer quality. By making repeated document use cheaper, LazyAttention can help long-context AI systems respond faster and use less GPU memory.