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 ক্যাশ বহিষ্কার অপ্টিমাইজেশন
বৃহৎ ভাষা মডেল (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 ক্যাশ উপাদান নির্বাচন করা।
ইনপুট: মোট বাজেট 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}
দীর্ঘ অনুক্রম প্রক্রিয়া বেঞ্চমার্ক (Ruler, LongBench ইত্যাদি)
সামগ্রিক মূল্যায়ন: এটি একটি উচ্চ-মানের গবেষণা পেপার যা তাত্ত্বিক অবদান এবং ব্যবহারিক মূল্যের মধ্যে ভাল ভারসাম্য অর্জন করে। Ada-KV পদ্ধতি সহজ এবং কার্যকর, তাত্ত্বিক বিশ্লেষণ কঠোর, পরীক্ষামূলক যাচাইকরণ ব্যাপক। পেপারটি শুধুমাত্র বিদ্যমান পদ্ধতির গুরুত্বপূর্ণ সীমাবদ্ধতা সমাধান করে না, বরং ভবিষ্যত গবেষণার জন্য মূল্যবান কাঠামো এবং দিকনির্দেশনা প্রদান করে।