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 индекс, который обеспечивает эффективное и точное извлечение критических токенов с помощью таких методов, как трипартитная аппроксимация внимания, оценка внимания с ограниченной точностью и сегментированная кластеризация. Это дополняется wave buffer, который координирует размещение кэша KV и перекрывает вычисления и передачу данных между GPU и CPU для поддержания высокой пропускной способности. В отличие от предыдущих методов, основанных на разреженности, которые сталкиваются с проблемами выбора токенов и координации оборудования, RetroInfer обеспечивает стабильную производительность без ущерба для точности модели. Эксперименты на бенчмарках с длинным контекстом показывают ускорение до 4,5 раз по сравнению с полным вниманием в пределах памяти GPU и до 10,5 раз по сравнению с базовыми методами разреженного внимания, когда кэш KV расширяется до памяти CPU, при этом сохраняя точность на уровне полного внимания.

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.