Learn from the Past: Fast Sparse Indexing for Large Language Model Decoding

📄 arXiv: 2506.15704v1 📥 PDF

作者: Feiyu Yao, Qian Wang

分类: cs.LG, cs.AI, cs.CL

发布日期: 2025-05-30

💡 一句话要点

LFPS：利用历史注意力模式加速长文本LLM解码中的稀疏索引检索

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 长文本解码 稀疏注意力 索引检索 历史注意力模式

📋 核心要点

  1. 长文本LLM解码时，KV缓存的内存需求是瓶颈，稀疏注意力虽能缓解，但索引计算开销大。
  2. LFPS利用历史注意力模式，动态构建稀疏索引候选，预测Top-k索引，降低索引检索成本。
  3. 实验表明，LFPS在长文本基准测试中，相对于完整注意力和精确Top-k检索，显著提升解码速度。

📝 摘要（中文）

随着大型语言模型(LLMs)支持的上下文长度不断增加，解码过程中键值(KV)缓存的内存需求迅速增长，成为GPU内存容量和PCIe带宽的关键瓶颈。稀疏注意力机制通过仅计算选定键值对的注意力权重来缓解此问题。然而，它们的索引计算通常需要遍历所有键向量，导致显著的计算和数据传输开销。为了降低索引检索的成本，现有方法通常将每个解码步骤视为独立过程，未能利用历史解码信息中嵌入的时间相关性。为此，我们提出LFPS(Learn From the Past for Sparse Indexing)，一种基于历史注意力模式动态构建稀疏索引候选的加速方法。LFPS捕捉解码器注意力中的两种普遍趋势——垂直模式（关注固定位置）和斜线模式（关注相对位置），并结合位置扩展策略，有效地预测当前步骤的Top-k索引。我们使用Llama-3.1-8B-Instruct作为基础模型，在具有挑战性的长上下文基准测试（如LongBench-RULER）上验证了LFPS。实验结果表明，在RTX 4090 GPU和一个Xeon Gold 6430的单个CPU核心上，LFPS实现了高达22.8倍的速度提升（相对于完整注意力）和9.6倍的速度提升（相对于精确的Top-k检索），同时保持了生成准确性。这些结果表明，LFPS为长上下文LLM推理中的解码优化提供了一种实用且高效的解决方案。

🔬 方法详解

问题定义：论文旨在解决长文本LLM解码过程中，由于KV缓存的内存需求过大，以及现有稀疏注意力机制索引计算开销高昂的问题。现有方法将每个解码步骤视为独立过程，忽略了历史解码信息中的时间相关性，导致效率低下。

核心思路：LFPS的核心思路是“从过去学习”，即利用历史解码过程中的注意力模式来预测当前解码步骤中需要关注的键值对。通过捕捉历史注意力中的垂直模式（关注固定位置）和斜线模式（关注相对位置），并结合位置扩展策略，可以有效地缩小索引检索范围，降低计算开销。

技术框架：LFPS方法主要包含以下几个阶段：1) 历史注意力模式分析：分析历史解码步骤中的注意力权重，识别垂直模式和斜线模式。2) 索引候选构建：基于识别出的注意力模式，构建当前解码步骤的索引候选集合。3) 位置扩展：对索引候选集合进行位置扩展，以覆盖更广泛的上下文信息。4) Top-k索引选择：从扩展后的索引候选集合中选择Top-k个索引，用于后续的注意力计算。

关键创新：LFPS的关键创新在于其动态构建稀疏索引候选的方法，该方法充分利用了历史解码信息中的时间相关性，避免了对所有键向量的遍历。与现有方法相比，LFPS能够更准确地预测当前解码步骤中需要关注的键值对，从而显著降低索引检索的计算开销。

关键设计：LFPS的关键设计包括：1) 垂直模式和斜线模式的识别算法，需要根据具体的注意力权重分布进行调整。2) 位置扩展策略，需要平衡扩展范围和计算开销。3) Top-k索引的选择标准，需要考虑生成质量和计算效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LFPS在LongBench-RULER等长文本基准测试中，使用Llama-3.1-8B-Instruct模型，在RTX 4090 GPU上实现了高达22.8倍的速度提升（相对于完整注意力），在Xeon Gold 6430 CPU上实现了9.6倍的速度提升（相对于精确的Top-k检索），同时保持了生成准确性。这些结果充分验证了LFPS在长文本LLM解码优化方面的有效性。

🎯 应用场景

LFPS可应用于各种需要处理长文本的LLM应用场景，例如长篇文档摘要、机器翻译、对话系统等。通过降低解码过程中的内存需求和计算开销，LFPS能够提升LLM的推理速度和效率，使其能够更好地服务于实际应用，并降低部署成本。未来，该方法有望进一步扩展到其他类型的序列模型和任务中。

📄 摘要（原文）

As large language models (LLMs) continue to support increasingly longer contexts, the memory demand for key-value (KV) caches during decoding grows rapidly, becoming a critical bottleneck in both GPU memory capacity and PCIe bandwidth. Sparse attention mechanisms alleviate this issue by computing attention weights only for selected key-value pairs. However, their indexing computation typically requires traversing all key vectors, resulting in significant computational and data transfer overhead. To reduce the cost of index retrieval, existing methods often treat each decoding step as an independent process, failing to exploit the temporal correlations embedded in historical decoding information. To this end, we propose LFPS(Learn From the Past for Sparse Indexing), an acceleration method that dynamically constructs sparse indexing candidates based on historical attention patterns. LFPS captures two prevalent trends in decoder attention -vertical patterns (attending to fixed positions) and slash patterns (attending to relative positions) -and incorporates a positional expansion strategy to effectively predict the Top-k indices for the current step. We validate LFPS on challenging long-context benchmarks such as LongBench-RULER, using Llama-3.1-8B-Instruct as the base model. Experimental results show that LFPS achieves up to 22.8$\times$ speedup over full attention and 9.6$\times$ speedup over exact Top-k retrieval on an RTX 4090 GPU and a single CPU core of a Xeon Gold 6430, respectively, while preserving generation accuracy. These results demonstrate that LFPS offers a practical and efficient solution for decoding optimization in long-context LLM inference.