中文

从过去学习:大语言模型解码的稀疏索引快速方法

机器学习 2025-06-23 v1 人工智能 计算与语言

摘要

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

关键词

引用

@article{arxiv.2506.15704,
  title  = {Learn from the Past: Fast Sparse Indexing for Large Language Model Decoding},
  author = {Feiyu Yao and Qian Wang},
  journal= {arXiv preprint arXiv:2506.15704},
  year   = {2025}
}