HeadInfer: Энергоэффективный вывод больших языковых моделей за счет поэтапного выгрузки голов

HeadInfer: Memory-Efficient LLM Inference by Head-wise Offloading

February 18, 2025

Авторы: Cheng Luo, Zefan Cai, Hanshi Sun, Jinqi Xiao, Bo Yuan, Wen Xiao, Junjie Hu, Jiawei Zhao, Beidi Chen, Anima Anandkumar

cs.AI

Аннотация

Трансформерные большие языковые модели (LLM) демонстрируют впечатляющие результаты в генерации длинных контекстов. Увеличение длины контекста непропорционально сместило объем памяти, занимаемый LLM во время вывода, в сторону кэша ключей и значений (KV cache). В данной работе мы предлагаем HEADINFER, который переносит KV cache в оперативную память CPU, избегая необходимости полного хранения KV cache для любого слоя трансформера на GPU. HEADINFER использует детализированную стратегию выгрузки по головам, сохраняя на GPU только KV cache для избранных голов внимания, при этом динамически вычисляя выход внимания. С помощью анализа roofline мы показываем, что HEADINFER сохраняет вычислительную эффективность, значительно сокращая объем занимаемой памяти. Мы оцениваем HEADINFER на модели Llama-3-8B с последовательностью из 1 миллиона токенов, сокращая объем памяти GPU, занимаемый KV cache, с 128 ГБ до 1 ГБ, а общее использование памяти GPU — с 207 ГБ до 17 ГБ, что соответствует снижению на 92% по сравнению с базовым выводом в формате BF16. Примечательно, что HEADINFER позволяет выполнять вывод для 4 миллионов токенов на модели с 8 миллиардами параметров на одном потребительском GPU с 24 ГБ памяти (например, NVIDIA RTX 4090) без использования методов аппроксимации.

English

Transformer-based large language models (LLMs) demonstrate impressive performance in long context generation. Extending the context length has disproportionately shifted the memory footprint of LLMs during inference to the key-value cache (KV cache). In this paper, we propose HEADINFER, which offloads the KV cache to CPU RAM while avoiding the need to fully store the KV cache for any transformer layer on the GPU. HEADINFER employs a fine-grained, head-wise offloading strategy, maintaining only selective attention heads KV cache on the GPU while computing attention output dynamically. Through roofline analysis, we demonstrate that HEADINFER maintains computational efficiency while significantly reducing memory footprint. We evaluate HEADINFER on the Llama-3-8B model with a 1-million-token sequence, reducing the GPU memory footprint of the KV cache from 128 GB to 1 GB and the total GPU memory usage from 207 GB to 17 GB, achieving a 92% reduction compared to BF16 baseline inference. Notably, HEADINFER enables 4-million-token inference with an 8B model on a single consumer GPU with 24GB memory (e.g., NVIDIA RTX 4090) without approximation methods.