RetroInfer: Un Enfoque de Almacenamiento Vectorial para la Inferencia Escalable de LLM con Contexto Largo

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

Resumen

El creciente tamaño de contexto de los modelos de lenguaje grandes (LLMs, por sus siglas en inglés) presenta desafíos significativos para la inferencia eficiente, principalmente debido a las limitaciones de memoria y ancho de banda de las GPU. Presentamos RetroInfer, un sistema novedoso que reconceptualiza la caché clave-valor (KV) como un sistema de almacenamiento de vectores que aprovecha la dispersión inherente de la atención para acelerar la inferencia de LLMs con contextos largos. En su núcleo se encuentra el índice wave, un índice vectorial consciente de la atención (Attention-aWare VEctor index) que permite la recuperación eficiente y precisa de tokens críticos mediante técnicas como la aproximación de atención tripartita, la estimación de atención con límite de precisión y el clustering segmentado. Esto se complementa con el buffer wave, que coordina la ubicación de la caché KV y superpone el cálculo y la transferencia de datos entre la GPU y la CPU para mantener un alto rendimiento. A diferencia de métodos previos basados en dispersión que luchan con la selección de tokens y la coordinación de hardware, RetroInfer ofrece un rendimiento robusto sin comprometer la precisión del modelo. Los experimentos en benchmarks de contexto largo muestran una aceleración de hasta 4.5X sobre la atención completa dentro de los límites de memoria de la GPU y hasta 10.5X sobre los baselines de atención dispersa cuando la caché KV se extiende a la memoria de la CPU, todo ello manteniendo una precisión equivalente a la de la atención 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.