ShadowKV : Cache KV dans les Ombres pour LLM à Long Context et Haut Débit Inference

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

October 28, 2024

Auteurs: Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, Beidi Chen

cs.AI

Résumé

Avec le déploiement généralisé des grands modèles de langage à contexte long (LLM), il y a une demande croissante pour un support efficace de l'inférence à haut débit. Cependant, à mesure que le cache clé-valeur (KV) s'étend avec la longueur de la séquence, l'empreinte mémoire croissante et la nécessité d'y accéder pour chaque génération de jeton entraînent une faible capacité de traitement lors du service des LLM à contexte long. Alors que diverses méthodes d'attention dynamique et éparses ont été proposées pour accélérer l'inférence tout en maintenant la qualité de génération, elles échouent soit à réduire de manière suffisante la consommation de mémoire GPU, soit introduisent une latence de décodage significative en déchargeant le cache KV vers le CPU. Nous présentons ShadowKV, un système d'inférence LLM à contexte long à haut débit qui stocke le cache clé de rang faible et décharge le cache de valeur pour réduire l'empreinte mémoire pour des tailles de lots plus importantes et des séquences plus longues. Pour minimiser la latence de décodage, ShadowKV utilise une stratégie précise de sélection KV qui reconstruit de manière minimale des paires KV éparses à la volée. En évaluant ShadowKV sur un large éventail de référentiels, y compris RULER, LongBench et Needle In A Haystack, et des modèles tels que Llama-3.1-8B, Llama-3-8B-1M, GLM-4-9B-1M, Yi-9B-200K, Phi-3-Mini-128K et Qwen2-7B-128K, nous démontrons qu'il peut prendre en charge des tailles de lots jusqu'à 6 fois plus grandes et augmenter le débit jusqu'à 3,04 fois sur un GPU A100 sans sacrifier la précision, dépassant même les performances réalisables avec une taille de lot infinie en supposant une mémoire GPU infinie. Le code est disponible sur 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.