Enregistré dans:
Détails bibliographiques
Auteurs principaux: Sharma, Akshat, Ding, Hangliang, Li, Jianping, Dani, Neel, Zhang, Minjia
Format: Preprint
Publié: 2024
Sujets:
Accès en ligne:https://arxiv.org/abs/2411.18077
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866912418527969280
author Sharma, Akshat
Ding, Hangliang
Li, Jianping
Dani, Neel
Zhang, Minjia
author_facet Sharma, Akshat
Ding, Hangliang
Li, Jianping
Dani, Neel
Zhang, Minjia
contents How to efficiently serve LLMs in practice has become exceptionally challenging due to their prohibitive memory and computation requirements. In this study, we investigate optimizing the KV cache, whose memory footprint poses a critical bottleneck in LLM inference, especially when dealing with long context tasks. To tackle the challenge, we introduce MiniKV, a KV cache optimization method that simultaneously preserves long context task accuracy while significantly reducing KV cache size via a novel 2-bit layer-discriminative KV cache. More importantly, we develop specialized CUDA kernels to make MiniKV compatible with FlashAttention. Experiments on a wide range of long context tasks show that MiniKV effectively achieves 86% KV cache compression ratio while recovering over 98.5% of accuracy, outperforming state-of-the-art methods while achieving excellent measured system performance improvements.
format Preprint
id arxiv_https___arxiv_org_abs_2411_18077
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle MiniKV: Pushing the Limits of LLM Inference via 2-Bit Layer-Discriminative KV Cache
Sharma, Akshat
Ding, Hangliang
Li, Jianping
Dani, Neel
Zhang, Minjia
Computation and Language
Machine Learning
How to efficiently serve LLMs in practice has become exceptionally challenging due to their prohibitive memory and computation requirements. In this study, we investigate optimizing the KV cache, whose memory footprint poses a critical bottleneck in LLM inference, especially when dealing with long context tasks. To tackle the challenge, we introduce MiniKV, a KV cache optimization method that simultaneously preserves long context task accuracy while significantly reducing KV cache size via a novel 2-bit layer-discriminative KV cache. More importantly, we develop specialized CUDA kernels to make MiniKV compatible with FlashAttention. Experiments on a wide range of long context tasks show that MiniKV effectively achieves 86% KV cache compression ratio while recovering over 98.5% of accuracy, outperforming state-of-the-art methods while achieving excellent measured system performance improvements.
title MiniKV: Pushing the Limits of LLM Inference via 2-Bit Layer-Discriminative KV Cache
topic Computation and Language
Machine Learning
url https://arxiv.org/abs/2411.18077