让缓存的注意力计算不绑定位置，提升长文本生成效率
LazyAttention: Efficient Retrieval-Augmented Generation with Deferred Positional Encoding
- 将位置编码延迟到计算时处理，实现无需复制的通用缓存复用
- 在倾斜文档分布下，首字生成时间快1.37倍，吞吐量提升40%
- 适合需要长上下文推理的场景，如检索增强生成和提示学习
键值（KV）缓存通过重用已生成词元的计算加速大语言模型推理，在检索增强生成（RAG）和上下文学习（ICL）等长上下文应用中尤为重要。然而，传统KV缓存直接将位置信息嵌入缓存，限制了其可复用性。现有方案要么仅允许前缀复用，要么需昂贵的内存重编码。我们提出LazyAttention，一种新型注意力机制，通过内核化延迟位置编码，实现零复制、与位置无关的KV复用。通过在注意力内核中动态调整位置编码，解决材料化瓶颈，使单一物理缓存副本可服务于任意位置的多个逻辑请求。采用针对预填充和解码优化的注意力内核，系统显著提升效率：在倾斜文档分布下，相比当前最优的Block-Attention，首字生成时间（TTFT）减少1.37倍，推理吞吐量提升1.40倍，同时保持相当的输出质量。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。