RetroInfer: A Vector-Storage Approach for Scalable Long-Context LLM Inference

📄 arXiv: 2505.02922v2 📥 PDF

作者: Yaoqi Chen, Jinkai Zhang, Baotong Lu, Qianxi Zhang, Chengruidong Zhang, Jingjia Luo, Di Liu, Huiqiang Jiang, Qi Chen, Jing Liu, Bailu Ding, Xiao Yan, Jiawei Jiang, Chen Chen, Mingxing Zhang, Yuqing Yang, Fan Yang, Mao Yang

分类: cs.LG

发布日期: 2025-05-05 (更新: 2025-06-30)

备注: 17 pages

💡 一句话要点

RetroInfer：一种向量存储方法，用于可扩展的长上下文LLM推理。

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

关键词: 长上下文LLM 推理加速 注意力稀疏性 向量存储 KV缓存

📋 核心要点

  1. 长上下文LLM推理面临GPU内存和带宽瓶颈，现有方法难以兼顾效率与精度。
  2. RetroInfer将KV缓存视为向量存储，利用注意力稀疏性，通过波索引和波缓冲区加速推理。
  3. 实验表明，RetroInfer在长上下文任务中显著优于现有方法，同时保持了模型精度。

📝 摘要（中文）

大型语言模型（LLM）不断增长的上下文长度对高效推理提出了重大挑战，这主要是由于GPU内存和带宽的限制。我们提出了RetroInfer，这是一种新颖的系统，它将键值（KV）缓存重新概念化为向量存储系统，该系统利用固有的注意力稀疏性来加速长上下文LLM推理。其核心是波索引，这是一种注意力感知向量索引，它通过三方注意力近似、精度有界的注意力估计和分段聚类等技术，能够高效准确地检索关键token。与之互补的是波缓冲区，它协调KV缓存的放置，并重叠GPU和CPU之间的计算和数据传输，以维持高吞吐量。与之前在token选择和硬件协调方面存在困难的基于稀疏性的方法不同，RetroInfer在不影响模型准确性的前提下，提供了强大的性能。在长上下文基准测试中进行的实验表明，在GPU内存限制内，速度比完全注意力机制提高了4.5倍，当KV缓存扩展到CPU内存时，速度比稀疏注意力基线提高了10.5倍，同时保持了完全注意力机制的准确性。

🔬 方法详解

问题定义：论文旨在解决长上下文LLM推理中，由于KV缓存导致的GPU内存和带宽瓶颈问题。现有方法，如稀疏注意力机制，在token选择和硬件协调方面存在困难，难以在效率和精度之间取得平衡。

核心思路：RetroInfer的核心思路是将KV缓存重新概念化为向量存储系统，并利用注意力机制的稀疏性，只保留和检索对推理至关重要的token。通过设计一种注意力感知的向量索引（wave index）和缓存管理机制（wave buffer），实现高效且准确的长上下文推理。

技术框架：RetroInfer主要包含两个核心模块：波索引（wave index）和波缓冲区（wave buffer）。波索引负责构建和维护KV缓存的向量索引，以便快速检索重要的token。波缓冲区负责协调KV缓存的放置，并重叠GPU和CPU之间的计算和数据传输，以提高吞吐量。整体流程包括：输入文本 -> 注意力估计 -> 波索引检索 -> KV缓存更新 -> 模型推理 -> 输出结果。

关键创新：RetroInfer的关键创新在于波索引的设计，它是一种注意力感知的向量索引，能够高效准确地检索关键token。与现有方法不同，波索引采用三方注意力近似、精度有界的注意力估计和分段聚类等技术，能够在保证精度的前提下，显著减少需要存储和计算的token数量。

关键设计：波索引的关键设计包括：1) 三方注意力近似，用于快速估计token的重要性；2) 精度有界的注意力估计，用于控制检索的精度；3) 分段聚类，用于减少索引的大小。波缓冲区的关键设计包括：1) KV缓存的动态放置策略，根据token的重要性选择放置在GPU或CPU内存中；2) 计算和数据传输的重叠，以提高吞吐量。具体的参数设置和损失函数等技术细节在论文中未详细说明，属于未知信息。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

RetroInfer在长上下文基准测试中表现出色，在GPU内存限制内，速度比完全注意力机制提高了4.5倍。当KV缓存扩展到CPU内存时，速度比稀疏注意力基线提高了10.5倍，同时保持了与完全注意力机制相当的准确性。这些结果表明，RetroInfer是一种高效且准确的长上下文LLM推理解决方案。

🎯 应用场景

RetroInfer可应用于需要处理长文本的各种场景，例如长篇文档摘要、代码生成、对话系统、知识库问答等。通过降低长上下文LLM推理的计算和存储成本，RetroInfer能够促进这些应用在资源受限环境中的部署，并提升用户体验。未来，该技术有望进一步扩展到其他类型的模型和任务中。

📄 摘要（原文）

The growing context lengths of large language models (LLMs) pose significant challenges for efficient inference, primarily due to GPU memory and bandwidth constraints. We present RetroInfer, a novel system that reconceptualizes the key-value (KV) cache as a vector storage system which exploits the inherent attention sparsity to accelerate long-context LLM inference. At its core is the wave index, an Attention-aWare VEctor index that enables efficient and accurate retrieval of critical tokens through techniques such as tripartite attention approximation, accuracy-bounded attention estimation, and segmented clustering. Complementing this is the wave buffer, which coordinates KV cache placement and overlaps computation and data transfer across GPU and CPU to sustain high throughput. Unlike prior sparsity-based methods that struggle with token selection and hardware coordination, RetroInfer delivers robust performance without compromising model accuracy. Experiments on long-context benchmarks show up to 4.5X speedup over full attention within GPU memory limits and up to 10.5X over sparse attention baselines when KV cache is extended to CPU memory, all while preserving full-attention-level accuracy.