LazyAttention: Efficient Retrieval-Augmented Generation with Deferred Positional Encoding

2026-06-02

AI Summary

Traditional key-value (KV) caching struggles to reuse cached states across different positions due to embedded positional information, limiting inference efficiency in long-context scenarios such as retrieval-augmented generation. This work proposes LazyAttention, a novel mechanism that dynamically defers positional encoding within the attention kernel, enabling for the first time zero-copy, position-agnostic KV cache reuse. By designing custom prefill and decode attention kernels integrated with delayed positional encoding, our approach overcomes the limitations of existing methods that either restrict reusable prefixes or incur high memory overhead. Under skewed document distributions, the system reduces time-to-first-token by 1.37× and improves inference throughput by 1.40× compared to Block-Attention, while maintaining comparable generation quality.

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.

Problem

• KV caching
• positional encoding
• retrieval-augmented generation
• cache reusability
• long-context inference

Innovation

• LazyAttention
• deferred positional encoding
• KV caching
• retrieval-augmented generation
• attention kernel