LazyAttention：通过延迟位置编码实现高效的检索增强生成

背景

大型语言模型（LLM）的推理加速在很大程度上依赖于键值（KV）缓存技术。KV 缓存通过复用过去计算的结果来加速生成 token 的过程，其重要性在长上下文应用场景中尤为突出，例如检索增强生成（RAG）和上下文学习（ICL）。在这些场景中，模型需要处理大量的历史上下文信息，KV 缓存的有效利用直接决定了系统的响应速度和资源消耗。

然而，传统的 KV 缓存机制存在一个根本性的局限：它将位置信息直接嵌入到缓存中。这意味着缓存中的 KV 向量是特定于某个绝对位置的。当不同的请求需要复用相同的上下文内容，但起始位置不同时，传统的缓存机制无法直接复用这些数据，因为位置编码不匹配。

现有的解决方案通常面临两难选择：要么限制复用范围，仅允许在固定的前缀（prefix）内进行复用，这极大地限制了灵活性；要么为了支持任意位置的重用，需要昂贵的内存显影（memory materialization）来进行位置重新编码，这带来了巨大的计算和内存开销。这种局限性阻碍了 KV 缓存在大规模、高并发 RAG 场景中的效率提升。

核心内容

为了解决上述问题，研究人员提出了 LazyAttention，这是一种新颖的注意力机制。其核心创新在于引入了“延迟位置编码”（Deferred Positional Encoding），并通过内核化（kernelization）技术实现了零拷贝（zero-copy）、与位置无关的 KV 复用。

LazyAttention 的工作原理主要包含以下几个关键方面：

1. 延迟位置编码与内核化： 传统方法在构建 KV 缓存时就确定了位置编码。LazyAttention 则不同，它在注意力计算内核（attention kernels）内部动态调整位置编码。这意味着 KV 缓存中存储的是“无位置”或“通用”的键值对，而具体的位置信息是在注意力计算发生的瞬间，通过内核代码动态注入的。

2. 零拷贝与位置无关的复用： 由于位置信息不再固化在缓存中，同一个物理 KV 副本可以被多个逻辑请求复用，无论这些请求在上下文中的起始位置如何。这实现了真正的“零拷贝”复用，极大地提高了内存利用率。

3. 针对预填充和解码的专用注意力内核： LazyAttention 系统针对预填充（prefilling，即处理输入提示词阶段）和解码（decoding，即生成 token 阶段）设计了专门的注意力内核。这种针对性优化确保了在保持高吞吐量的同时，能够高效地处理动态的位置调整。

4. 解决显影瓶颈： 通过上述机制，LazyAttention 成功解决了位置重新编码所需的内存显影瓶颈。它允许单个物理 KV 副本服务于多个逻辑请求，且这些请求可以位于任意位置，无需为每个请求创建独立的 KV 副本或进行昂贵的实时编码计算。

关键要点

• 技术突破：LazyAttention 是一种基于延迟位置编码的新型注意力机制，实现了 KV 缓存的零拷贝、位置无关复用。
• 解决痛点：克服了传统 KV 缓存因嵌入固定位置信息而导致的复用局限性，以及现有方案中昂贵的内存显影开销。
• 核心机制：通过在注意力内核中动态调整位置编码，而非在缓存构建阶段固化位置信息，使得单一物理 KV 副本可服务于多个不同起始位置的逻辑请求。
• 性能提升：在偏斜文档分布（skewed document distributions）的场景下，与最先进的 Block-Attention 方法相比：
  • 首 token 延迟（TTFT）降低了 1.37倍（即速度提升约 37%）。
  • 推理吞吐量提高了 1.40倍。
• 质量保持：在显著提升效率的同时，LazyAttention 保持了与基线方法相当输出质量，未因优化而牺牲模型性能。

意义与影响

LazyAttention 的提出对长上下文 LLM 应用，特别是 RAG 系统，具有重要的工程意义。

首先，它显著提升了资源利用效率。在 RAG 场景中，大量文档片段可能被多个查询复用。LazyAttention 允许这些片段在内存中只保留一份副本，通过逻辑映射服务于不同的查询，从而大幅降低内存占用，使得在有限硬件资源下处理更长的上下文或更多的并发请求成为可能。

其次，它优化了推理延迟。通过减少内存显影和拷贝操作，LazyAttention 降低了首 token 延迟（TTFT），这对于用户体验至关重要的交互式应用而言是一个关键改进。同时，吞吐量的提升意味着单位时间内可以处理更多的请求，降低了运营成本。

最后，LazyAttention 为注意力机制的设计提供了新的思路。它证明了通过软件层面的内核优化（如延迟计算特定元数据）可以有效突破硬件缓存机制的限制，这种“软硬化结合”的思路可能在未来其他 AI 系统优化中得到广泛应用。