RetroInfer: 大規模長文脈LLM推論のためのベクトルストレージアプローチ

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

May 5, 2025

著者: 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.AI

要旨

大規模言語モデル（LLM）のコンテキスト長の増加は、主にGPUメモリと帯域幅の制約により、効率的な推論に大きな課題をもたらしています。本論文では、RetroInferという新しいシステムを提案します。このシステムは、キー・バリュー（KV）キャッシュをベクトルストレージシステムとして再概念化し、内在するアテンションのスパース性を活用して長文コンテキストのLLM推論を加速します。その中核となるのは、wave index（Attention-aWare VEctor index）です。このインデックスは、三分近似アテンション、精度保証型アテンション推定、セグメント化クラスタリングなどの技術を通じて、重要なトークンの効率的かつ正確な検索を可能にします。これを補完するのがwave bufferで、KVキャッシュの配置を調整し、GPUとCPU間の計算とデータ転送をオーバーラップさせて高いスループットを維持します。トークン選択とハードウェア調整に苦戦する従来のスパース性ベースの手法とは異なり、RetroInferはモデルの精度を損なうことなく堅牢な性能を提供します。長文コンテキストベンチマークでの実験では、GPUメモリ制限内でフルアテンションと比較して最大4.5倍の高速化を実現し、KVキャッシュをCPUメモリに拡張した場合にはスパースアテンションベースラインと比較して最大10.5倍の高速化を達成しました。これらはすべて、フルアテンションレベルの精度を維持しながら行われています。

English

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.