RetroInfer: Uma Abordagem de Armazenamento Vetorial para Inferência Escalável de LLM com Contexto Longo

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

May 5, 2025

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

Resumo

O aumento dos comprimentos de contexto em grandes modelos de linguagem (LLMs) apresenta desafios significativos para a inferência eficiente, principalmente devido às limitações de memória e largura de banda da GPU. Apresentamos o RetroInfer, um sistema inovador que reconceitualiza o cache de chave-valor (KV) como um sistema de armazenamento de vetores que explora a esparsidade inerente da atenção para acelerar a inferência de LLMs com contextos longos. No seu núcleo está o índice wave, um índice vetorial consciente da atenção (Attention-aWare VEctor index) que permite a recuperação eficiente e precisa de tokens críticos por meio de técnicas como aproximação tripartida da atenção, estimação de atenção com limite de precisão e clustering segmentado. Complementando isso está o buffer wave, que coordena o posicionamento do cache KV e sobrepõe computação e transferência de dados entre GPU e CPU para manter uma alta taxa de transferência. Diferente de métodos baseados em esparsidade anteriores que lutam com a seleção de tokens e a coordenação de hardware, o RetroInfer oferece desempenho robusto sem comprometer a precisão do modelo. Experimentos em benchmarks de contexto longo mostram uma aceleração de até 4,5X em relação à atenção completa dentro dos limites de memória da GPU e até 10,5X em relação às linhas de base de atenção esparsa quando o cache KV é estendido para a memória da CPU, tudo isso mantendo a precisão no nível da atenção completa.

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.