ملخص تقني: RDKV - تخصيص معدل التشويه لبتات ذاكرة الـ KV لعمليات الاستبعاد والكمية المشتركة لذاكرة التخزين المؤقت KV

1. بيان المشكلة

تواجه النماذج اللغوية الكبيرة (LLMs) اختناقًا حرجًا أثناء الاستدلال في السياقات الطويلة: عرض نطاق الذاكرة وسعتها. ينمو حجم ذاكرة التخزين المؤقت للمفاتيح والقيم (KV cache) خطيًا مع طول التسلسل، مما يتطلب إعادة قراءة عشرات الجيجابايت من البيانات من ذاكرة النطاق العالي (HBM) خارج الشريحة إلى الذاكرة داخل الشريحة عند كل خطوة فك تشفير. وهذا يجعل عملية الاستدلال مقيدة بالذاكرة بدلاً من أن تكون مقيدة بالحوسبة.

تندرج استراتيجيات الضغط الحالية عمومًا ضمن فئتين معزولتين:

1. الاستبعاد (Eviction): تقييم الرموز (tokens) أو القنوات (channels) والتخلص من الأقل تصنيفًا بالكامل (إما الاحتفاظ أو الحذف الثنائي). هذا لا يترك أي مستوى دقة للرموز ذات الأهمية المتوسطة.
2. الكمية (Quantization): الاحتفاظ بجميع الرموز مع تقليل عرض البتات الخاص بها. هذا يحتفظ حتى بالرموز الضئيلة، مما يهدر الذاكرة.

بينما تحاول بعض الحلول الهجينة الحديثة الجمع بينهما، إلا أنها غالبًا ما تعتمد على درجات توجيه تجريبية أو نسب مخصصة مسبقًا يتم تحديدها عبر بحث خارجي (offline search)، مما يفشل في التعامل مع الاستبعاد والكمية كمسألة تحسين موحدة.

2. المنهجية: RDKV

يقترح البحث إطار عمل RDKV (ضغط ذاكرة KV بمعدل التشويه)، وهو إطار عمل يصيغ ضغط ذاكرة KV كمسألة معدل-تشويه (rate-distortion). ومن هذا المنظور، لا يعد الاستبعاد والكمية إجراءين منفصلين، بل هما طرفا منحنى واحد لتخصيص البتات.

2.1 الصيغة النظرية

يستنتج RDKV وزن التشويه لكل وحدة تخزين مؤقت (رمز في ذاكرة V، وقناة في ذاكرة K) بناءً على الاضطراب الناجم عن إزالتها أو كميتها على عملية حساب الانتباه.

• ذاكرة V (الرموز): وزن التشويه $w_t$ للرمز هو مسافة التباين الإجمالي لتوزيع الانتباه بعد الاستبعاد. ويُحسب كمجموع درجات الانتباه المستلمة من قبل الرمز عبر جميع مواضع الاستعلام: $w_t = \sum_{\tau} a_{\tau,t}$.
• ذاكرة K (القنوات): وزن التشويه $w_c$ للقناة هو النمط الطيفي لانحراف اللوجيت (logit deviation) الناتج عن تصفير تلك القناة. ويُحسب كحاصل ضرب معايير أعمدة الاستعلام والمفتاح المقابلة: $w_c = \frac{1}{\sqrt{d}} \|Q_{:,c}\|_2 \|K_{:,c}\|_2$.

تعمل هذه الأوزان كمعاملات تكلفة في دالة هدف معدل-التشويه:
$$\min \sum_{u} w_u \sigma_u 2^{-b_u} \quad \text{s.t.} \quad \sum_{u} b_u \leq B$$
حيث $b_u$ هو عرض البت المخصص للوحدة $u$، و $B$ هو إجمالي ميزانية البتات.

2.2 خط أنابيب تخصيص البتات

يتم حل التخصيص عبر تعبئة المياه العكسية (reverse water-filling) (الاسترخاء اللاجرانجي) المكيفة لعروض البتات المنفصلة المدعومة عتاديًا $\mathcal{B} = \{0, 2, 4, 8, 16\}$.

1. المرحلة 1 (الوزن): حساب أوزان الرموز ($w_t$) وأوزان القنوات ($w_c$) مرة واحدة بعد مرحلة التعبئة (prefill) باستخدام ذاكرة التخزين المؤقت الكاملة غير المضغوطة.
2. المرحلة 2 (تخصيص الرموز): حل مسألة الحقيبة متعددة الخيارات (MCKP) لذاكرة V. يتم تعيين عروض بت للرموز من المجموعة $\{0, 2, 4, 8, 16\}$. عرض البت 0 يقابل الاستبعاد.
3. المرحلة 3 (تخصيص القنوات): حل مسألة MCKP لذاكرة K، ولكن فقط للرموز التي تم الاحتفاظ بها في المرحلة 2. يتم تعيين عروض البت للقنوات بشكل مستقل عن عرض بت الـ V الخاص بالرمز، مما يسمح للرمز الواحد بامتلاك مكونات K و V ذات دقة مختلطة.

2.3 فك التشفيد المجمع الفعال (TriZone)

لتحقيق وفورات فعلية في الذاكرة وزمن الوصول، يقدم RDKV تخطيط ذاكرة التخزين المؤقت TriZone الذي يدمج عملية إزالة الكمية (dequantization) داخل نواة الانتباه (attention kernel)، متجنبًا الحاجة إلى فك حزم البيانات إلى FP16 قبل المعالجة.

• المنطقة A (Zone A): تدفق الذاكرة القديمة المجمع الذي يحتوي على صفوف K و V المحتفظ بها مع عروض البت $\{2, 4, 8\}$. يتم تجميعها حسب عرض البت للسماح بمسارات إزالة كمية موحدة.
• المنطقة B (Zone B): صفوف V بصيغة FP16 للرموز المخصصة لدقة 16 بت.
• المنطقة C (Zone C): صفوف K و V بصيغة FP16 للرموز المولدة حديثًا.

تقوم نواة الانتباه بحساب اللوجيت والمخرجات عبر دمج هذه المناطق الثلاث، مع دمج عملية إزالة الكمية مباشرة في عملية ضرب المصفوفات، مما يلغي عبء إضافي على عرض نطاق الذاكرة.

3. المساهمات الرئيسية

1. صياغة معدل-التشويه: صاغ الباحثون ضغط ذاكرة KV كمسألة معدل-تشويه موحدة، مثبتين أن الحل الأمثل يمزج بطبيعته بين الوحدات ذات المعدل الصفري (الاستبعاد) والوحدات ذات المعدل المحدود (الكمية). هذا يلغي الحاجة إلى النسب التجريبية أو عمليات البحث الخارجية.
2. التخصيص المنفصل عبر تعبئة المياه العكسية: قام المؤلفون بتطبيق التخصيب كمسألة حقيبة منفصلة يمكن حلها عبر الاسترخاء اللاجرانجي، وتعيين عروض البت المثلى من $\{0, 2, 4, 8, 16\}$ بناءً على أوزان التشويه المستمدة من ديناميكيات الانتباه.
3. تخطيط TriZone المجمع: تنفيذ مدرك للعتاد يدمج إزالة الكمية في نواة الانتباه، محولًا وفورات عرض البت النظرية إلى تقليلات فعلية في ذروة الذاكرة وزمن وصول فك التشفيد.

4. النتائج التجريبية

أُجريت التجارب على نماذج LLaMA-3.1-8B-Instruct، و Mistral-7B، و Qwen3-4B، و LLaMA-2-13B، و Qwen2.5-72B عبر اختبارات LongBench، و RULER، و InfiniteBench.

• الدقة: يتفوق RDKV على أفضل خط أساس تم تقييمه (Ada-SnapKV) بنسبة 9.1% في المتوسط. في LongBench، يستعيد 97.81% من دقة الذاكرة الكاملة مع الاحتفاظ بـ 2.48% فقط من الذاكرة.
• أداء السياق الطويل:
  • RULER: عند سياق 128K، يحقق RDKV دقة 66.95% مقارنة بـ 58.43% لأقوى خط أساس (AdaKV)، بفارق 8.5 نقطة.
  • InfiniteBench: يحتل RDKV المرتبة الأولى بين طرق الضغط بمتوسط درجة 39.46، متفوقًا على AdaKV بمقدار 1.40 نقطة.
  • الإبرة في كوم القش (Needle-in-a-Haystack): يحافظ RDKV على دقة استرجاع شبه موحدة (متوسط 0.99) عبر الأعماق، بينما تفشل خطوط الأساس التي تعتمد على الاستبعاد فقط بشكل كبير عند الأعماق المتوسطة.
• الكفاءة:
  • زمن الوصول (Latency): مع سياق 128K، يحقق RDKV تسريعًا في فك التشفيد بمقدار 4.5× مقارنة بفك تشفير FlashAttention-2 للذاكرة الكاملة.
  • الذاكرة: يحقق تقليلاً في ذروة الذاكرة بمقدار 1.9× عند سياق 128K، وينجح في تشغيل سياق 256K على وحدة GPU واحدة من نوع A100 64GB، حيث يفشل FullKV بسبب نفاذ الذاية (OOM).

5. الأهمية والادعاءات

يزعم البحث أن التعامل مع الاستبعاد والكمية معًا ضمن إطار تحسين موحد هو أفضل من النهج المرحلي أو القائم على النسب. من خلال اشتقاق درجات التخصيص مباشرة من التشويه الناجم عن حساب الانتباه، يتجنب RDKV "الفخ الثنائي" لطرق الاستبعاد فقط (التي تتخلص من الرموز متوسطة الأهمية) و"الهدر" في طرق الكمية فقط (التي تحتفظ بالرموت الضئيلة).

يؤكد المؤلفون أن طريقتهم:

• لا تتطلب إعادة تدريب النموذج.
• محايدة تجاه البنية (architecture-agnostic)، وتعمل باستمرار عبر أحجام وعائلات النماذج المختلفة.
• توفر مقايضة مثالية (Pareto-optimal) بين زمن الوصول والدقة، وتتفوق بصرامة على خطوط الأساس في كلا المقياسين عند أطوال السياق الطويلة.
• تثبت أن "الأمثل المستمر" لنظرية معدل-التشويه، عندما يتم جعلها منفصلة للعتاد، ينتج طبيعيًا استراتيجية بتات مختلطة تتفوق على أحدث التقنيات الحالية.

يخلص البحث إلى أن هذا النهج المشترك يحسن الدقة دون التضحية بالكفاءة، مما يتيح الاستدلال بسياق طويل على الأجهزة محدودة الموارد.