2025-11-16T09:58:12.370377

Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference

Feng, Lv, Cao et al.

Large Language Models have excelled in various domains but face efficiency challenges due to the growing Key-Value (KV) cache required for long-sequence inference. Recent efforts aim to reduce KV cache size by evicting vast non-critical cache elements during runtime while preserving generation quality. However, these methods typically allocate compression budgets uniformly across all attention heads, ignoring the unique attention patterns of each head. In this paper, we establish a theoretical loss upper bound between pre- and post-eviction attention output, explaining the optimization target of prior cache eviction methods, while guiding the optimization of adaptive budget allocation. Base on this, we propose {\it Ada-KV}, the first head-wise adaptive budget allocation strategy. It offers plug-and-play benefits, enabling seamless integration with prior cache eviction methods. Extensive evaluations on 13 datasets from Ruler and 16 datasets from LongBench, all conducted under both question-aware and question-agnostic scenarios, demonstrate substantial quality improvements over existing methods. Our code is available at https://github.com/FFY0/AdaKV.

academic

Ada-KV: कुशल LLM अनुमान के लिए अनुकूली बजट आवंटन द्वारा KV कैश निष्कासन को अनुकूलित करना

मूल जानकारी

पेपर ID: 2407.11550
शीर्षक: Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference
लेखक: Yuan Feng, Junlin Lv, Yukun Cao, Xike Xie, S. Kevin Zhou
वर्गीकरण: cs.CL cs.AI
प्रकाशन समय/सम्मेलन: तंत्रिका सूचना प्रसंस्करण प्रणाली पर 39वां सम्मेलन (NeurIPS 2025)
पेपर लिंक: https://arxiv.org/abs/2407.11550

सारांश

बड़े भाषा मॉडल (LLMs) विभिन्न क्षेत्रों में उत्कृष्ट प्रदर्शन करते हैं, लेकिन लंबे अनुक्रम अनुमान में बढ़ती Key-Value (KV) कैश आवश्यकताओं के कारण दक्षता चुनौतियों का सामना करते हैं। हाल के शोध रनटाइम पर गैर-महत्वपूर्ण कैश तत्वों को निष्कासित करके KV कैश आकार को कम करते हैं, जबकि पीढ़ी की गुणवत्ता बनाए रखते हैं। हालांकि, ये विधियां आमतौर पर सभी ध्यान सिर के बीच संपीड़न बजट को समान रूप से वितरित करती हैं, प्रत्येक सिर के अद्वितीय ध्यान पैटर्न को नजरअंदाज करती हैं। यह पेपर निष्कासन से पहले और बाद में ध्यान आउटपुट के बीच सैद्धांतिक हानि की ऊपरी सीमा स्थापित करता है, पिछली कैश निष्कासन विधियों के अनुकूलन उद्देश्यों की व्याख्या करता है, जबकि अनुकूली बजट आवंटन के अनुकूलन को निर्देशित करता है। इसके आधार पर, लेखकों ने Ada-KV प्रस्तावित किया, पहली सिर-स्तरीय अनुकूली बजट आवंटन रणनीति। यह विधि प्लग-एंड-प्ले लाभ प्रदान करती है, जो मौजूदा कैश निष्कासन विधियों के साथ निर्बाध रूप से एकीकृत हो सकती है।

अनुसंधान पृष्ठभूमि और प्रेरणा

समस्या विवरण

जैसे-जैसे बड़े भाषा मॉडल अनुक्रम लंबाई को संसाधित करते हैं (जैसे GPT 128K का समर्थन करता है, Claude3 200K का समर्थन करता है, Gemini-Pro-1.5 2M टोकन का समर्थन करता है), KV कैश की मेमोरी आवश्यकता तेजी से बढ़ती है। 8B पैरामीटर वाले LLM के लिए, एक एकल 2M टोकन अनुक्रम को संसाधित करने के लिए 256GB तक कैश की आवश्यकता हो सकती है, जो GPU मेमोरी दक्षता और कम्प्यूटेशनल रनटाइम दक्षता को गंभीर रूप से प्रभावित करती है।

मौजूदा विधियों की सीमाएं

मौजूदा कैश निष्कासन विधियां मुख्य रूप से दो श्रेणियों में विभाजित हैं:

स्लाइडिंग विंडो निष्कासन विधियां: सरलता से प्रारंभिक और हाल के कैश तत्वों को बनाए रखते हैं, लेकिन पीढ़ी की गुणवत्ता में काफी कमी लाते हैं
Top-k निष्कासन विधियां: ध्यान वजन के आधार पर महत्वपूर्ण कैश तत्वों का चयन करते हैं, लेकिन सभी ध्यान सिर के बीच बजट को समान रूप से वितरित करते हैं

मुख्य समस्या यह है कि मौजूदा विधियां विभिन्न ध्यान सिर की अद्वितीय विशेषताओं को नजरअंदाज करती हैं: कुछ सिर में विरल ध्यान केंद्रित पैटर्न होते हैं, जबकि अन्य सिर का ध्यान वितरण अधिक बिखरा हुआ होता है।

अनुसंधान प्रेरणा

Llama-3.1-8B-Instruct मॉडल का विश्लेषण करके, लेखकों ने पाया कि अधिकांश ध्यान सिर केवल कैश के एक छोटे अनुपात (जैसे शीर्ष 5%) के साथ लगभग सभी ध्यान वजन को बनाए रख सकते हैं, जबकि बिखरे हुए सिर को अधिक बड़े कैश अनुपात की आवश्यकता होती है। यह असमान ध्यान केंद्रण पैटर्न अनुकूली बजट आवंटन के लिए सैद्धांतिक आधार प्रदान करता है।

मुख्य योगदान

अनुकूली बजट आवंटन रणनीति: पहली सिर-स्तरीय अनुकूली बजट आवंटन रणनीति Ada-KV प्रस्तावित करता है, जो प्रत्येक ध्यान सिर के अद्वितीय ध्यान पैटर्न के अनुसार गतिशील रूप से बजट आवंटन को समायोजित कर सकता है
सैद्धांतिक ढांचा स्थापना: कैश निष्कासन के लिए एक सैद्धांतिक ढांचा स्थापित करता है, निष्कासन हानि को परिभाषित करता है और इसकी ऊपरी सीमा प्राप्त करता है, मौजूदा विधियों के अनुकूलन उद्देश्यों की व्याख्या करता है और Ada-KV के डिजाइन को निर्देशित करता है
प्लग-एंड-प्ले संगतता: Ada-KV में प्लग-एंड-प्ले विशेषता है, जो मौजूदा कैश निष्कासन विधियों में निर्बाध रूप से एकीकृत हो सकता है, और कुशल CUDA कर्नेल कार्यान्वयन के माध्यम से कम्प्यूटेशनल दक्षता बनाए रखता है
व्यापक प्रायोगिक सत्यापन: Ruler और LongBench के 29 डेटासेट पर व्यापक मूल्यांकन किया गया है, समस्या-जागरूक और समस्या-अज्ञेय दोनों परिदृश्यों में महत्वपूर्ण सुधार दिखाता है

विधि विवरण

कार्य परिभाषा

एक बहु-सिर स्व-ध्यान परत को देखते हुए, बजट बाधा के तहत KV कैश तत्वों को बनाए रखने के लिए चुनें, जिससे निष्कासन के बाद ध्यान आउटपुट और मूल आउटपुट के बीच हानि को कम किया जा सके।

सैद्धांतिक आधार

L1 निष्कासन हानि परिभाषा

लेखकों ने निष्कासन हानि को निष्कासन से पहले और बाद में स्व-ध्यान तंत्र आउटपुट के बीच L1 दूरी के रूप में परिमाणित किया:

$\text{L1 Eviction Loss} = ||y - \hat{y}||_1$

जहां $y$ और $\hat{y}$ क्रमशः निष्कासन से पहले और बाद में ध्यान आउटपुट हैं।

हानि ऊपरी सीमा व्युत्पत्ति

प्रमेय 3.1: L1 निष्कासन हानि को $\epsilon$ ऊपरी सीमा द्वारा सीमित किया जा सकता है:

$\text{L1 Eviction Loss} \leq \epsilon = 2hC - 2C\sum_{i \in [1,h]}\sum_{j \in [1,n]} I_i^j A_i^j$

जहां $C = \max\{\|V_iW_i^O\|_\infty\}$ एक स्थिरांक है, $I_i^j$ निष्कासन निर्णय संकेतक चर है, $A_i^j$ ध्यान वजन है।

प्रमेय 3.2: Top-k कैश निष्कासन विधि दिए गए बजट आवंटन के तहत हानि ऊपरी सीमा को कम कर सकती है:

$\epsilon^* = 2hC - 2C\sum_{i \in [1,h]}\sum_{A_i^j \in \text{Top-k}(A_i, k=B_i)} A_i^j$

Ada-KV एल्गोरिदम

एल्गोरिदम 1: अनुकूली बजट आवंटन

इनपुट: कुल बजट B, प्रत्येक सिर ध्यान वजन {A_i}
आउटपुट: आवंटित बजट {B_i^*}
1. सभी सिर के ध्यान वजन को जोड़ें: A = Cat({A_i})
2. A से शीर्ष B वजन का चयन करें: Top-k(A, k=B)
3. प्रत्येक सिर के लिए चयनित वजन की संख्या गिनें: {f_i}
4. आवंटित बजट सेट करें: {B_i^* = f_i}

सैद्धांतिक लाभ

प्रमेय 3.3: अनुकूली बजट आवंटन न्यूनतम हानि ऊपरी सीमा को प्राप्त कर सकता है:

$\epsilon^{**} = \min_{\{B_i\}} \epsilon^*$

मौजूदा विधियों के साथ एकीकरण

लेखकों ने Ada-KV के दो SOTA विधियों के साथ एकीकरण को प्रदर्शित किया:

Ada-SnapKV और Ada-Pyramid

एल्गोरिदम 2 के माध्यम से, Ada-KV को SnapKV और Pyramid में निर्बाध रूप से एकीकृत किया जा सकता है:

अवलोकन विंडो के भीतर ध्यान वजन की गणना करें
Ada-KV एल्गोरिदम का उपयोग करके बजट आवंटित करें
अत्यधिक विरल आवंटन को रोकने के लिए सुरक्षा पैरामीटर α = 0.2 लागू करें
Top-k निष्कासन निर्णय निष्पादित करें

तकनीकी नवाचार बिंदु

वैश्विक अनुकूलन दृष्टिकोण: सिर-स्तरीय बजट आवंटन को स्थानीय अनुकूलन के बजाय वैश्विक अनुकूलन समस्या के रूप में देखता है
सैद्धांतिक निर्देशित डिजाइन: कठोर सैद्धांतिक विश्लेषण के आधार पर एल्गोरिदम डिजाइन को निर्देशित करता है
कम्प्यूटेशनल दक्षता गारंटी: परिवर्तनशील-लंबाई FlashAttention और समतल कैश लेआउट के माध्यम से कम्प्यूटेशनल दक्षता बनाए रखता है
GQA संगतता: Group Query Attention का समर्थन करता है, अतिरिक्त कैश संपीड़न को प्राप्त करता है

प्रायोगिक सेटअप

डेटासेट

Ruler बेंचमार्क: 13 लंबे अनुक्रम कार्य, मुख्य रूप से Needle-in-a-Haystack परीक्षण के वेरिएंट, 16K लंबाई का मूल्यांकन
LongBench बेंचमार्क: 16 डेटासेट, एकल-दस्तावेज़ QA, बहु-दस्तावेज़ QA, सारांश, कम-नमूना शिक्षा, सिंथेटिक कार्य और कोड पीढ़ी को कवर करते हैं

आधार मॉडल

Llama-3.1-8B-Instruct
Mistral-7B-instruct-v0.2

मूल्यांकन मेट्रिक्स

कार्य प्रकार के अनुसार संबंधित मेट्रिक्स का उपयोग करें: F1 स्कोर (QA कार्य), Rouge-L (सारांश कार्य), सटीकता (वर्गीकरण कार्य), संपादन समानता (कोड कार्य)

तुलना विधियां

आधार विधियां: SnapKV, Pyramid, StreamingLLM
संवर्धित संस्करण: Ada-SnapKV, Ada-Pyramid

प्रायोगिक परिदृश्य

समस्या-जागरूक संपीड़न: समस्या ज्ञात मानक परिदृश्य
समस्या-अज्ञेय संपीड़न: अधिक चुनौतीपूर्ण वास्तविक अनुप्रयोग परिदृश्य

प्रायोगिक परिणाम

मुख्य परिणाम

Ruler बेंचमार्क परीक्षण

समस्या-अज्ञेय परिदृश्य में, Llama-3.1-8B-Instruct का उपयोग करके:

80% कैश बजट: Ada-SnapKV ने SnapKV के स्कोर को 87.59 से 92.67 तक बढ़ाया
20% कैश बजट: Ada-SnapKV ने SnapKV के स्कोर को 44.02 से 53.29 तक बढ़ाया

LongBench बेंचमार्क परीक्षण

समस्या-अज्ञेय परिदृश्य में:

Ada-SnapKV और Ada-Pyramid सभी निश्चित बजट सेटिंग्स के तहत लगातार पीढ़ी की गुणवत्ता में सुधार करते हैं
2048 बजट के तहत लगभग नुकसान-मुक्त प्रदर्शन के करीब

उप-कार्य विश्लेषण

कठिन Needle-in-a-Haystack कार्य में:

S-NIAH-3 कार्य (80% बजट): Ada-SnapKV ने SnapKV को 62.4 से 97.6 तक बढ़ाया
MK-NIAH-2 कार्य (80% बजट): Ada-SnapKV ने SnapKV को 85.2 से 99.6 तक बढ़ाया

कम्प्यूटेशनल दक्षता

Ada-SnapKV निश्चित 1024 बजट के तहत:

पीक मेमोरी उपयोग मूल SnapKV के समान है
डिकोडिंग विलंबता मूल SnapKV के समान है
दोनों पूर्ण कैश स्थिति की तुलना में काफी बेहतर हैं

व्यापक अनुप्रयोग सत्यापन

Ada-KV रणनीति को कई बाद के कार्यों द्वारा अपनाया गया है:

CriticalKV + Ada-KV: 20% कैश के तहत 42.99 से 43.77 तक
DefensiveKV + Ada-KV: 20% कैश के तहत 43.78 से 46.68 तक

निष्कर्ष और चर्चा

मुख्य निष्कर्ष

Ada-KV पहली बार सिर-स्तरीय अनुकूली बजट आवंटन रणनीति प्रस्तावित करता है, जो मौजूदा कैश निष्कासन विधियों के प्रदर्शन में काफी सुधार करता है
सैद्धांतिक विश्लेषण कैश निष्कासन के लिए एक कठोर ढांचा स्थापित करता है, एल्गोरिदम डिजाइन को निर्देशित करता है
समस्या-अज्ञेय संपीड़न परिदृश्य मौजूदा विधियों की सीमाओं को प्रकट करता है, जिसे अधिक ध्यान दिया जाना चाहिए

सीमाएं

वर्तमान सिर-स्तरीय आवंटन अभी भी एकल परत तक सीमित है, परत-पार आवंटन तक विस्तारित नहीं है
सुरक्षा पैरामीटर α को विभिन्न बजटों के तहत प्रदर्शन को संतुलित करने की आवश्यकता है
सैद्धांतिक विश्लेषण L1 दूरी पर आधारित है, वास्तविक पीढ़ी की गुणवत्ता को पूरी तरह से प्रतिबिंबित नहीं कर सकता है

भविष्य की दिशाएं

सिर-स्तरीय आवंटन तंत्र को परत-पार परिदृश्य तक विस्तारित करें
संबंधित परत-पार सैद्धांतिक विश्लेषण विकसित करें
प्रशिक्षण-समय सिर महत्व विश्लेषण के साथ संयोजित करें
अन्य अनुकूलन तकनीकों (जैसे परिमाणीकरण, विरल ध्यान) के साथ संयुक्त अनुकूलन

गहन मूल्यांकन

लाभ

ठोस सैद्धांतिक योगदान: एक पूर्ण सैद्धांतिक ढांचा स्थापित करता है, हानि ऊपरी सीमा व्युत्पत्ति से एल्गोरिदम डिजाइन तक तर्क स्पष्ट है
सरल और प्रभावी विधि: एल्गोरिदम सरल और समझने में आसान है, प्लग-एंड-प्ले विशेषता इसे अपनाने में आसान बनाती है
व्यापक और पर्याप्त प्रयोग: 29 डेटासेट पर व्यापक मूल्यांकन, अनदेखी समस्या-अज्ञेय परिदृश्य को शामिल करता है
उच्च व्यावहारिक मूल्य: कई बाद के कार्यों द्वारा अपनाया गया है, विधि के मूल्य और प्रभाव को प्रमाणित करता है

कमियां

सिद्धांत और अभ्यास के बीच अंतराल: हालांकि सैद्धांतिक रूप से हानि ऊपरी सीमा को कम किया गया है, लेकिन वास्तविक हानि को कम करने की गारंटी नहीं दे सकता है
हाइपरपैरामीटर संवेदनशीलता: सुरक्षा पैरामीटर α का चयन अनुभवजन्य ट्यूनिंग की आवश्यकता है
विस्तार सीमाएं: वर्तमान में केवल एकल परत के भीतर बजट आवंटन पर विचार करता है
मूल्यांकन सीमाएं: मुख्य रूप से मध्यम आकार के मॉडल पर मूल्यांकन, बड़े पैमाने के मॉडल पर प्रभाव सत्यापन की प्रतीक्षा में है

प्रभाव

शैक्षणिक योगदान: KV कैश अनुकूलन क्षेत्र के लिए नई अनुसंधान दिशा प्रदान करता है
व्यावहारिक मूल्य: प्लग-एंड-प्ले विशेषता इसे वास्तविक सिस्टम में तैनात करने में आसान बनाती है
पुनरुत्पादनीयता: खुला स्रोत कोड और विस्तृत कार्यान्वयन विवरण प्रदान करता है
प्रेरणादायक: बाद के अनुसंधान के लिए सैद्धांतिक ढांचा और पद्धति मार्गदर्शन प्रदान करता है

लागू परिदृश्य

लंबे अनुक्रम अनुमान: विशेष रूप से लंबे संदर्भ को संसाधित करने की आवश्यकता वाले अनुप्रयोगों के लिए उपयुक्त
संसाधन-सीमित वातावरण: GPU मेमोरी सीमित होने की स्थिति में अनुमान दक्षता को अनुकूलित करता है
वास्तविक समय प्रणाली: गुणवत्ता और दक्षता को संतुलित करने की आवश्यकता वाली ऑनलाइन सेवाएं
बहु-मोड़ संवाद: समस्या-अज्ञेय संपीड़न परिदृश्य विशेष रूप से संवाद प्रणालियों के लिए उपयुक्त है

संदर्भ

पेपर में 64 संबंधित संदर्भों का हवाला दिया गया है, मुख्य रूप से शामिल हैं:

बड़े भाषा मॉडल मौलिक कार्य (GPT-4, Claude, Gemini आदि)
KV कैश अनुकूलन विधियां (H2O, SnapKV, Pyramid आदि)
ध्यान तंत्र अनुकूलन (FlashAttention, विरल ध्यान आदि)
लंबे अनुक्रम प्रसंस्करण बेंचमार्क (Ruler, LongBench आदि)

समग्र मूल्यांकन: यह एक उच्च गुणवत्ता का अनुसंधान पेपर है जो सैद्धांतिक योगदान और व्यावहारिक मूल्य के बीच अच्छा संतुलन प्राप्त करता है। Ada-KV विधि सरल लेकिन प्रभावी है, सैद्धांतिक विश्लेषण कठोर है, प्रायोगिक सत्यापन पर्याप्त है। पेपर न केवल मौजूदा विधियों की महत्वपूर्ण सीमाओं को हल करता है, बल्कि भविष्य के अनुसंधान के लिए मूल्यवान ढांचा और दिशा भी प्रदान करता है।