ShadowKV: KV-кэш в тенях для высокопроизводительного LLM-вывода с длинным контекстом

ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference

October 28, 2024

Авторы: Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, Beidi Chen

cs.AI

Аннотация

С широким развертыванием моделей большого размера с длинным контекстом (LLM) возросла потребность в эффективной поддержке высокопроизводительного вывода. Однако по мере расширения кеша ключ-значение (KV) с увеличением длины последовательности увеличивается объем памяти и необходимость обращения к нему для каждой генерации токена, что приводит к низкой производительности при обслуживании моделей LLM с длинным контекстом. Хотя были предложены различные методы динамического разреженного внимания для ускорения вывода при сохранении качества генерации, они либо не удается достаточно снизить потребление памяти GPU, либо вносят значительную задержку декодирования путем выгрузки кеша KV на ЦП. Мы представляем ShadowKV, систему вывода моделей LLM с длинным контекстом высокой производительности, которая хранит кеш ключей низкого ранга и выгружает кеш значений для снижения объема памяти при больших размерах пакетов и длинных последовательностях. Для минимизации задержки декодирования ShadowKV использует точную стратегию выбора KV, которая восстанавливает минимальные разреженные пары KV на лету. Оценивая ShadowKV на широком спектре бенчмарков, включая RULER, LongBench и Needle In A Haystack, а также модели, такие как Llama-3.1-8B, Llama-3-8B-1M, GLM-4-9B-1M, Yi-9B-200K, Phi-3-Mini-128K и Qwen2-7B-128K, мы демонстрируем, что она может поддерживать до 6 раз большие размеры пакетов и увеличивать производительность до 3.04 раз на GPU A100 без ущерба точности, превосходя даже производительность, достижимую с бесконечным размером пакета при предположении о бесконечной памяти GPU. Код доступен по адресу https://github.com/bytedance/ShadowKV.

English

With the widespread deployment of long-context large language models (LLMs), there has been a growing demand for efficient support of high-throughput inference. However, as the key-value (KV) cache expands with the sequence length, the increasing memory footprint and the need to access it for each token generation both result in low throughput when serving long-context LLMs. While various dynamic sparse attention methods have been proposed to speed up inference while maintaining generation quality, they either fail to sufficiently reduce GPU memory consumption or introduce significant decoding latency by offloading the KV cache to the CPU. We present ShadowKV, a high-throughput long-context LLM inference system that stores the low-rank key cache and offloads the value cache to reduce the memory footprint for larger batch sizes and longer sequences. To minimize decoding latency, ShadowKV employs an accurate KV selection strategy that reconstructs minimal sparse KV pairs on-the-fly. By evaluating ShadowKV on a broad range of benchmarks, including RULER, LongBench, and Needle In A Haystack, and models like Llama-3.1-8B, Llama-3-8B-1M, GLM-4-9B-1M, Yi-9B-200K, Phi-3-Mini-128K, and Qwen2-7B-128K, we demonstrate that it can support up to 6times larger batch sizes and boost throughput by up to 3.04times on an A100 GPU without sacrificing accuracy, even surpassing the performance achievable with infinite batch size under the assumption of infinite GPU memory. The code is available at https://github.com/bytedance/ShadowKV.