RetroInfer : Une Approche de Stockage Vectoriel pour l'Inférence Évolutive des LLM à Contexte Long

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

May 5, 2025

Auteurs: 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

Résumé

L'augmentation des longueurs de contexte des grands modèles de langage (LLMs) présente des défis majeurs pour l'inférence efficace, principalement en raison des contraintes de mémoire GPU et de bande passante. Nous présentons RetroInfer, un système novateur qui reconceptualise le cache clé-valeur (KV) en un système de stockage vectoriel exploitant la sparsité intrinsèque de l'attention pour accélérer l'inférence des LLMs à contexte long. Au cœur de ce système se trouve l'index wave, un index Attention-aWare VEctor qui permet une récupération efficace et précise des tokens critiques grâce à des techniques telles que l'approximation d'attention tripartite, l'estimation d'attention à précision bornée et le clustering segmenté. Ceci est complété par le wave buffer, qui coordonne le placement du cache KV et superpose le calcul et le transfert de données entre le GPU et le CPU pour maintenir un débit élevé. Contrairement aux méthodes basées sur la sparsité antérieures qui peinent avec la sélection des tokens et la coordination matérielle, RetroInfer offre des performances robustes sans compromettre la précision du modèle. Les expériences sur des benchmarks à contexte long montrent une accélération allant jusqu'à 4,5X par rapport à l'attention complète dans les limites de mémoire GPU, et jusqu'à 10,5X par rapport aux bases de référence d'attention sparse lorsque le cache KV est étendu à la mémoire CPU, tout en préservant une précision équivalente à celle de l'attention complète.

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.