SIFT：利用注意力不变性实现RAG预填充的快速选择性索引计算

检索增强生成 (RAG) 将 LLM 查询与相关文档一起注入，以提高响应质量。此注入增加了提示长度并减慢了第一个令牌 (TTFT) 的时间。与标准查询不同，RAG 查询具有上下文重用的独特属性，其中相同的文档在用户查询中重复出现。因此，为每个 RAG 查询完全重新计算文档会产生冗余计算并增加 TTFT。之前的工作离线预计算 RAG 文档的 KV 张量，并在在线预填充期间粗略地重新计算一些标记。然而，由于高延迟的磁盘传输，这种 KV 重用通常比现代 GPU 上的完全重新计算慢。此外，这种粗粒度的重新计算降低了精度。为了解决这些限制，本文提出了 SIFT：Selective-Index For Fast Compute of RAG Prefill by Exploiting Attention Invariance。 SIFT 离线处理文档并为每个文档提取高注意力分数的细粒度位置。接下来，我们确定以下注意力不变性见解，使我们能够在运行时利用提取的位置：（1）局部注意力不变性：文档中高注意力分数的位置与周围文档保持不变。这有助于我们预测文档关注的高分位置。 （2）交叉注意力一致性：文档内注意力较高的键也会吸引后续文档的交叉注意力。这有助于我们预测该文档涉及未来文档的高分位置。重要的是，SIFT 不存储 KV 数据，仅以两个紧凑位向量的形式存储高分位置。 SIFT 的存储空间比 KV 张量小 24,000 倍，从而避免了昂贵的磁盘传输。在预填充期间，SIFT 仅计算标记位置的注意力，并将 TTFT 提高 1.71 倍，同时将完全重新计算的精度保持在 1% 以内。