LazyAttention: Efficient Retrieval-Augmented Generation with Deferred Positional Encoding
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Question Answering | 2WikiMQA | -- | 66 | |
| LLM Inference Efficiency | Uniform Sampled Documents | TTFT (ms)191.7 | 9 | |
| Question Answering | RAG QA 5 retrieved documents | TTFT (ms)191.7 | 6 | |
| Question Answering | TriviaQA | Exact Match (EM)73 | 5 | |
| Question Answering | NarrativeQA | Exact Match (EM)59.7 | 5 | |
| Few-shot classification | AG News Few-shot | TTFT Speedup1.31 | 4 | |
| KV Cache Efficiency | Trace-driven simulation 1 GB KV cache budget 100k-doc pool Zipf doc popularity | Low Skewness Hit Ratio3.47 | 4 | |
| KV Cache Efficiency | Trace-driven simulation 5 GB KV cache budget 100k-doc pool Zipf doc popularity | Hit Ratio (Low Skewness)10.85 | 4 | |
| KV Cache Efficiency | Trace-driven simulation 10 GB KV cache budget 100k-doc pool Zipf doc popularity | Low Skewness Hit Ratio13.78 | 4 | |
| KV Cache Efficiency | Trace-driven simulation 50 GB KV cache budget 100k-doc pool Zipf doc popularity | Hit Ratio (Low Skewness)20.22 | 4 |