LazyAttention: Efficient Retrieval-Augmented Generation with Deferred Positional Encoding

2026-06-03Computation and Language

Computation and LanguageMachine Learning
AI summary

The authors describe a new method called LazyAttention that makes language models faster by improving how they reuse past information during text generation. Unlike older methods that tie position information directly to memory, LazyAttention delays adding position details until needed, so the same memory can be reused in different contexts without extra copying. This approach speeds up the start of generating text and boosts overall efficiency without lowering output quality. Their experiments show it outperforms previous techniques, especially with long documents.

Key-value cachingLarge language modelsPositional encodingAttention mechanismInference efficiencyRetrieval-augmented generationIn-context learningKernelizationMemory materialization
Authors
Haocheng Xia, Mihir Pamnani, Hanxi Fang, Supawit Chockchowwat, Yongjoo Park
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$\times$ and increases inference throughput by 1.40$\times$ compared to the state-of-the-art Block-Attention, while maintaining comparable output quality.