LazyAttention: Efficient Retrieval-Augmented Generation with Deferred Positional Encoding
作者: Haocheng Xia, Mihir Pamnani, Hanxi Fang, Supawit Chockchowwat, Yongjoo Park
分类: cs.CL, cs.LG
发布日期: 2026-06-03
备注: ICML 2026
💡 一句话要点
提出LazyAttention以解决长上下文生成中的KV缓存问题
🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 长上下文生成 关键值缓存 延迟位置编码 推理效率 自然语言处理
📋 核心要点
- 现有的KV缓存方法在长上下文生成中存在位置编码重用性差的问题,限制了效率。
- LazyAttention通过延迟位置编码和动态调整注意力核,解决了传统方法的物化瓶颈。
- 实验结果显示,LazyAttention在TTFT和推理吞吐量上分别提高了1.37倍和1.40倍,且输出质量保持稳定。
📝 摘要(中文)
关键值(KV)缓存通过重用过去的计算加速大型语言模型(LLMs)的推理,尤其在检索增强生成(RAG)和上下文学习(ICL)等长上下文应用中更为重要。然而,传统的KV缓存将位置信息直接嵌入缓存中,限制了其重用性。现有解决方案要么限制重用范围,要么需要昂贵的内存物化进行位置重新编码。本文提出LazyAttention,一种新颖的注意力机制,通过延迟位置编码实现零拷贝、位置无关的KV重用。LazyAttention通过动态调整注意力核中的位置编码,解决了物化瓶颈,使得单个物理KV副本可以服务于多个逻辑请求。实验表明,在文档分布不均的情况下,该系统将首次生成时间(TTFT)减少了1.37倍,推理吞吐量提高了1.40倍,同时保持了输出质量的可比性。
🔬 方法详解
问题定义:本文旨在解决长上下文生成中KV缓存的重用性不足问题。传统方法将位置信息嵌入缓存,导致重用受限,且需要昂贵的内存物化进行位置重新编码。
核心思路:LazyAttention的核心思路是通过延迟位置编码实现零拷贝的KV重用。通过在注意力核中动态调整位置编码,LazyAttention能够在不增加内存开销的情况下,支持多个逻辑请求的处理。
技术框架:LazyAttention的整体架构包括两个主要模块:注意力核和KV缓存。注意力核负责动态调整位置编码,而KV缓存则存储生成的中间结果。系统通过调整注意力核的参数,实现对不同位置的灵活支持。
关键创新:LazyAttention的最大创新在于其延迟位置编码机制,使得单个物理KV副本可以在多个逻辑请求中复用,显著提高了推理效率。这一设计与传统方法的直接嵌入方式形成了本质区别。
关键设计:在设计中,LazyAttention采用了特定的参数设置以优化注意力核的性能,确保在动态调整位置编码时不会影响输出质量。同时,网络结构经过精心设计,以支持高效的KV缓存管理和位置编码调整。
📊 实验亮点
LazyAttention在实验中表现出色,在文档分布不均的情况下,首次生成时间(TTFT)减少了1.37倍,推理吞吐量提高了1.40倍,相较于现有的Block-Attention方法,展现了显著的性能提升,同时保持了输出质量的稳定性。
🎯 应用场景
LazyAttention的研究成果在长上下文生成任务中具有广泛的应用潜力,尤其是在自然语言处理、对话系统和信息检索等领域。通过提高推理效率,该方法能够显著提升用户体验,并为实时应用提供支持。未来,LazyAttention可能会影响更多基于大规模语言模型的应用场景,推动相关技术的发展。
📄 摘要(原文)
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.