Zachary، خبير خوارزميات في Novita AI
مجال البحث: تسريع الاستدلال
مقدمة
على مدار العام الماضي، منذ ظهور H2O، كان هناك انتشار كبير للأوراق البحثية حول KV Sparsity. ومع ذلك، فإن أحد التحديات الكبيرة في التطبيقات الواقعية هو الفجوة الكبيرة بين البحث الأكاديمي والتنفيذ العملي. على سبيل المثال، تتبنى أطر العمل مثل vLLM تقنية PagedAttention، وهي تقنية ذاكرة قائمة على التقسيم الصفحي غير متوافقة مع معظم خوارزميات التشتت، أو حيث تؤدي هذه الخوارزميات بشكل أسوأ من PagedAttention. تمنع مثل هذه المشكلات خوارزميات التشتت من الاستخدام الفعال في بيئات الإنتاج.
من خلال مراجعة الأوراق الأكاديمية الحديثة من العام الماضي في مجال KV Sparsity، قمنا بإجراء تعديلات على إطار عمل vLLM بناءً على KV Sparsity. تماشيت هذه التعديلات مع ميزات التحسين المتأصلة في إطار عمل vLLM، مثل Continuous Batching و FlashAttention و PagedAttention. في النهاية، من خلال مقارنة الإطار المعدل بأطر الاستدلال الأحدث (SOTA)، باستخدام النماذج والمعايير والأجهزة الأكثر شيوعًا في الإنتاج، حققنا تسريعًا بنسبة 1.5x في أداء الاستدلال.
قبل مناقشة KV Sparsity، من الضروري معالجة خاصية Massive Activations في نماذج LLM. في LLMs، يكون عدد صغير من قيم التنشيط أكثر نشاطًا بشكل كبير من غيرها، وأحيانًا تتجاوز القيم الأقل نشاطًا بعامل 100,000 أو أكثر. بمعنى آخر، تلعب مجموعة فرعية صغيرة من التوكنات دورًا حاسمًا. تشكل هذه الملاحظة الأساس لتطبيق تقنيات KV Sparsity، حيث أن الاحتفاظ فقط بالتوكنات المهمة يمكن أن يعزز أداء الاستدلال.
تدعم بيانات التحليل المحددة من Llama2 و Llama3 هذه الملاحظة بشكل أكبر، كما هو موضح في الشكل 1 والشكل 2.
الشكل 1. مقادير التنشيط (المحور z) في LLaMA2-7B. يمثل المحوران x و y أبعاد التسلسل والميزات. بالنسبة لهذا النموذج المحدد، نلاحظ أن التنشيطات ذات المقادير الهائلة تظهر في بُعدين ميزة ثابتين (1415، 2533)، ونوعين من التوكنات - توكن البداية، وأول نقطة (.) أو توكن السطر الجديد (\n).
على عكس Llama2، لا يُظهر Llama3 خاصية Massive Activations عبر جميع الطبقات. في الطبقات السفلية، يتم توزيع التنشيطات بشكل أكثر اتساقًا، بينما تعرض الطبقات الوسطى أنماط انتباه محلية. فقط في الطبقات العليا تصبح ظاهرة Massive Activations بارزة. تشرح هذه الملاحظة لماذا تتبنى نماذج مثل Llama3 تقنيات التشتت الهرمي، التي تُخصص التشتت لطبقات مختلفة وفقًا لأنماط تنشيطها المحددة.
الشكل 2. أنماط الانتباه للتوليد المعزز بالاسترجاع عبر الطبقات في LlaMa (Touvron et al., 2023a;b) تكشف أنه في الطبقات السفلية، يُظهر النموذج وضعًا واسع الطيف للانتباه، موزعًا درجات الانتباه بالتساوي على جميع المحتوى. في الطبقات الوسطى، يصبح الانتباه أكثر محلية داخل كل مستند، مما يشير إلى تجميع معلومات مُحسَّن (أشكال مثلثة حمراء منقطة في الطبقات 6 و 10). يبلغ هذا ذروته في الطبقات العليا، حيث يركز “الانتباه الهائل” على عدد قليل من التوكنات الرئيسية (أشرطة انتباه مركزة بعد الطبقة 18)، مستخرجًا المعلومات الأساسية للإجابات بكفاءة.
نموذج النظام
بعد تقديم خاصية Massive Activations في LLMs، أصبح لدينا فهم أساسي لمبادئ KV Sparsity. بعد ذلك، سنناقش كيفية تنفيذ KV Sparsity. في عملية استدلال النماذج الكبيرة، تتم مواجهة ثلاث اختناقات رئيسية عادةً: سعة ذاكرة GPU، والقدرة الحاسوبية، وعرض النطاق الترددي للإدخال/الإخراج. لتجنب الحسابات الزائدة عن الحاجة، غالبًا ما يتم تخزين KV مؤقتًا أثناء الاستدلال، مما يؤدي إلى استهلاك كبير لذاكرة GPU. وفي الوقت نفسه، يصبح الحمل الحسابي وإنتاجية الإدخال/الإخراج اختناقات حاسمة في مرحلتي التعبئة المسبقة (Prefill) وفك التشفير (Decode) على التوالي.
يعتمد تعديلنا لإطار عمل vLLM على خاصية Massive Activations في LLMs، مما ينفذ تشتتًا على مستوى الطبقة (أي مستويات مختلفة من التشتت لطبقات مختلفة). يقلل هذا النهج بشكل كبير من الحمل المرتبط بالتخزين المؤقت لـ KV، وبالتالي تسريع عملية استدلال نماذج LLM.
كما هو موضح في الشكل 3، أثناء عملية الاستدلال، نطبق KV Sparsity على طبقات مختلفة من النموذج. باستخدام استراتيجية تقليم، نقوم بإزالة أزواج KV ذات الدرجات المنخفضة مع الاحتفاظ بتلك ذات الدرجات الأعلى والقرب الأكبر. يقلل هذا النهج من استخدام الذاكرة، بالإضافة إلى الحمل الحسابي والإدخال/الإخراج، مما يؤدي في النهاية إلى تسريع الاستدلال.
الشكل 3. تأثير KV Sparsity على النموذج
نتائج تسريع الاستدلال
الاهتمام الرئيسي لمعظم الباحثين هو التأثير الفعلي لـ KV Sparsity في إطار عمل vLLM. قبل الخوض في التفاصيل، نقدم أولاً نتائج تقييم الأداء.
1. أداء الاستدلال
في تطبيقات LLM الواقعية، يعتبر طول الإدخال/الإخراج 4000/500 التكوين الأكثر شيوعًا. بالإضافة إلى ذلك، تُعد بطاقة RTX 4090 GPU واحدة من أكثر البطاقات استخدامًا في الصناعة. بناءً على هذه الشروط، أجرينا اختبارات مقارنة الأداء على دفعات. المجموعة الضابطة هي vLLM 0.6.1.p2، بينما المجموعة التجريبية هي Novita AI Sparse 0.5.1 (إصدار معدّل من vLLM 0.5.1 مع KV Sparsity). تم إجراء عدة جولات من اختبارات مقارنة الأداء بينهما، مع مقاييس رئيسية تشمل TTFT (الوقت حتى أول توكن، وهو أمر بالغ الأهمية لتجربة المستخدم) والإنتاجية (Throughput) (سرعة الاستدلال الفعلية).
تظهر نتائج الاختبار النهائية في الجدول 1. من خلال تطبيق KV Sparsity، تمكنا من زيادة إنتاجية vLLM بحوالي 1.58x، مع الحفاظ على TTFT قابل للاستخدام (P50 أقل من ثانية واحدة). بالإضافة إلى ذلك، في سيناريوهات حجم الدفعة الأكبر، وصل vLLM 0.6.1.p2 إلى حده الأقصى للأداء عند مستوى التزامن 10، بينما تمكن Novita AI Sparse 0.5.1 من الحفاظ على أداء TTFT مستقر حتى عند مستوى التزامن 20. يوضح هذا متانة KV Sparsity في ضمان استقرار الأداء في بيئات الإنتاج.
الجدول 1: مقارنة الأداء بين vLLM 0.6.1.p2 و Novita AI Sparse 0.5.1
2. دقة النموذج
نظرًا لأن KV Sparsity عبارة عن خوارزمية ضغط مع فقدان، فمن المهم بنفس القدر تقييم أداء النموذج. أجرينا تقييمات الأداء بناءً على Llama3–8B، بما في ذلك التقييمات باستخدام MMLU ومعايير أخرى. تشير النتائج إلى أن فقدان الدقة يظل في حدود حوالي 3%.
كما هو موضح في الجدول 2، اختبرنا أداء النموذج باستخدام معياري MMLU والعلوم الإنسانية:
الجدول 2: نتائج تقييم أداء Llama3–8B باستخدام معياري MMLU والعلوم الإنسانية
ثانيًا، أجرينا تقييمات لمهام الأسئلة والأجوبة في سيناريوهات النصوص الطويلة الشائعة في التطبيقات العملية، بأطوال إدخال تتراوح من 7k إلى 30k. تم الحفاظ على نسبة الضغط بأكثر من 10 مرات، كما هو موضح في الجدول 3. أشارت النتائج إلى فقدان إجمالي في الدقة بنسبة 10% تقريبًا.
الجدول 3: نتائج تقييم مهام الأسئلة والأجوبة للنصوص الطويلة بأطوال إدخال 7k إلى 30k
التقنيات الرئيسية
1. التشتت على مستوى الطبقة والتوازي الموتر
يستخدم إطار عمل vLLM استراتيجية جدولة على مستوى التكرار (Iteration-Level Scheduling)، والتي يشار إليها عادةً باسم Continuous Batching. تتميز طريقة الجدولة هذه بعدم انتظار كل تسلسل في الدفعة لإكمال التوليد؛ بدلاً من ذلك، تسهل الجدولة على مستوى التكرار، مما يؤدي إلى استخدام أعلى لوحدة معالجة الرسوميات (GPU) مقارنة بالتجميع الثابت. كما ذكرنا في بداية هذه الورقة، تمتلك نماذج مثل Llama3–8B بطبيعتها ميزات مثل التشتت الهرمي. وبالتالي، واجهت تعديلاتنا بناءً على Continuous Batching سلسلة من التحديات كما هو موضح في الشكل 4.
الشكل 4. إكمال سبعة تسلسلات باستخدام التجميع المستمر. على اليسار، ترى الدفعة بعد تكرار واحد، بينما على اليمين، الدفعة بعد عدة تكرارات. بمجرد أن يصدر التسلسل توكن نهاية التسلسل، نقدم تسلسلاً جديدًا في موضعه (على سبيل المثال، التسلسلات S5 و S6 و S7). تعزز هذه الطريقة استخدام GPU من خلال تجنب حاجة GPU لانتظار انتهاء جميع التسلسلات قبل بدء تسلسل جديد.
كان التحدي الرئيسي الذي عالجنا هو إدارة الذاكرة عبر الطبقات المتنوعة. نظرًا لاستخدام vLLM لنموذج إدارة ذاكرة صفحي موحد، واقتصار استراتيجية الطابور على دعم إما Full Cache أو Sliding Window، قمنا بتعديل البنية الأساسية لـ vLLM لتمكين دعم كلا الخيارين.
يمنح هذا التعديل لكل طبقة المرونة لاختيار إما Full Cache أو Sliding Window، مما يسمح بمستويات متميزة من التشتت عبر الطبقات، كما هو موضح في الرسم البياني المقدم.
التحدي الأول الذي يجب معالجته هو مشكلة إدارة الذاكرة عبر الطبقات المختلفة. نظرًا لأن vLLM يستخدم نموذج إدارة ذاكرة صفحي موحد لجميع الطبقات، وتسمح استراتيجية الطابور إما بـ Full Cache أو Sliding Window، نحتاج إلى تعديل البنية الأساسية لـ vLLM. على وجه التحديد، يتضمن ذلك تمكين الدعم لكل من Full Cache و Sliding Window بحيث يمكن للطبقات المختلفة اختيار استراتيجية إدارة الذاكرة المفضلة لديها بشكل مستقل. يسهل هذا التعديل في النهاية درجات متفاوتة من التشتت عبر الطبقات المختلفة. يظهر الرسم البياني الهيكلي المقابل في الشكل 5.
الشكل 5. هيكل إدارة الذاكرة على مستوى الطبقة في vLLM
بشكل عام، يمكن تصنيف التعديلات إلى ثلاث مراحل متميزة:
مرحلة تهيئة الخدمة: أثناء عملية بدء تشغيل الخدمة، تتم تهيئة الذاكرة المخصصة للطبقات الثلاث بشكل أساسي من خلال حساب المساحة المتبقية، إلى جانب تهيئة هرمية لهيكل البيانات الوصفية (Metadata). في المقابل، يعمل vLLM الأصلي كـ Full Cache كامل، دون مراعاة منطق تخصيص الذاكرة للطبقات المختلفة.
مرحلة التحضير للخطوة: أثناء عملية تهيئة الخطوات في vLLM، وتحديدًا في المراحل التي تتضمن أساليب مدير الكتلة (block manager) ومنفذ النموذج (model runner)، يتم استخدام استراتيجيات إدارة جدول كتلة مختلفة للطبقات المختلفة للتحكم في آلية تحديث KV. كما هو موضح في الشكل السابق، نستخدم آلية Full Cache للطبقات الثلاث الأولى، نخزين جميع KVs، بينما تستخدم الطبقات بعد الثالثة استراتيجية Sliding Window لاستبدال KVs ذات الدرجات الأدنى تلقائيًا. يسمح هذا النهج بتخصيص مرن لمساحة الذاكرة مصممة خصيصًا لاحتياجات كل طبقة، مما يوفر المساحة مع زيادة فعالية خوارزمية التشتت إلى أقصى حد.
مرحلة حساب التشتت (Sparsification Computation): تحدث هذه المرحلة عادةً بعد تنفيذ النموذج، حيث يتم تنفيذ عمليات التشتت بناءً على درجات التشتت المحسوبة. أحد الاعتبارات الرئيسية في هذه الخطوة هو التوافق مع تقنيات مثل Tensor Parallelism و CUDA graphs و GQA. في التطبيقات الهندسية العملية، يتطلب Tensor Parallelism اهتمامًا خاصًا. كما هو موضح في الشكل أدناه، في سيناريو بطاقة واحدة متعددة وحدات معالجة الرسوميات (single-machine multi-GPU)، من الضروري استخدام آلية تقسيم مناسبة لحصر عمليات تفتيت KV داخل كل GPU، وبالتالي تجنب الاتصال بين وحدات معالجة الرسوميات. لحسن الحظ، تستخدم وحدة الانتباه (Attention unit) نهج تقسيم قائم على ColumnLinear، مما يجعل تنفيذها بسيطًا نسبيًا.
الشكل 6. عملية من ثلاث مراحل لإدارة الذاكرة والتشتت في vLLM
تركز المناقشة أعلاه بشكل أساسي على التعديلات التي تم إجراؤها على مستوى الإطار لدعم التشتت الهرمي في vLLM. بعد ذلك، سنحول انتباهنا إلى التحسينات التي تم تنفيذها على مستوى CUDA.
2. تعديل الانتباه
تركز التعديلات على مستوى CUDA بشكل أساسي على التحسينات التي تم إجراؤها على وحدة حساب الانتباه، وتحديدًا فيما يتعلق بـ FlashAttention و PagedAttention. بينما يعتبر PagedAttention بسيطًا نسبيًا، سنؤكد على التعديلات المتعلقة بـ FlashAttention.
أولاً، نقدم صيغة حساب الانتباه:
كما هو موضح في الصيغة، تتضمن عملية حساب softmax حساب حاصل ضرب QK، وهي خوارزمية ذات مسارين (تتطلب تمريرين). ومع ذلك، يمكن تحقيق هدف الحساب النهائي O باستخدام خوارزمية FlashAttention بمسار واحد (تتطلب تمريرًا واحدًا). يمكن توضيح منطق تنفيذ FlashAttention من خلال الشكل 7 من ورقة FlashAttention2. باختصار، يتضمن تجزئة Q إلى أجزاء (chunks) لإجراء حسابات الكتل على KV، مع تحديث O و P وبيانات rowmax تدريجيًا حتى انتهاء الحلقة. أخيرًا، قسمة O على ℓ تمكن من حساب FlashAttention بمسار واحد.
الشكل 7. منطق تنفيذ FlashAttention2
بالنسبة لتنفيذنا لـ KV Sparsity، من المهم ملاحظة أنه أثناء عملية حساب FlashAttention، يتم بالفعل حساب القيم المطلوبة لـ softmax - وهي rowmax/ℓ (الحد الأقصى العالمي والقيم التراكمية). ومع ذلك، لا يُرجع FlashAttention نتائج هاتين القيمتين.
مع فهم أساسي لـ FlashAttention، يمكننا الآن مناقشة صيغ التسجيل للتشتت المقدمة في أوراق مثل PYRAMIDKV:
حيث [n − α, n] هو نطاق توكنات التعليمات. في كل طبقة l ولكل رأس h، يتم تحديد أفضل توكنات k l بأعلى الدرجات، ويتم الاحتفاظ بمخابئ KV الخاصة بها. يتم تجاهل جميع مخابئ KV الأخرى ولن يتم استخدامها في أي حسابات لاحقة طوال عملية التوليد.
بشكل أكثر تحديدًا، نحسب درجة اختيار التوكن i للاحتفاظ به في مخبأ KV كـ S في كل رأس انتباه h بواسطة:
في الصيغة أعلاه، يمثل A درجة softmax، بينما يمثل S مجموع درجات أقرب قيم softmax. كما هو موضح، softmax نفسه عبارة عن خوارزمية ذات مسارين، لا يمكن دمجها مباشرةً في FlashAttention. يعني هذا أيضًا أنه لا يمكن تنفيذ تسجيل التشتت (sparsification scoring) في مسار واحد.
كما ذكرنا سابقًا، على الرغم من أن FlashAttention لا يخرج درجات softmax، إلا أنه يحسب القيم الوسيطة مثل rowmax/ℓ. لذلك، من خلال الرجوع إلى Online Softmax: خوارزمية ذات مسارين، يمكننا استخدام مخرجات rowmax/ℓ لتحقيق التسجيل المحدد S لـ KV Sparsity. عمليًا، نحتاج فقط إلى إجراء تحسينات على FlashAttention الموضح في الشكل 7 لتلبية هذا المطلب.
الشكل 7. تعديل FlashAttention لتسجيل KV Sparsity
- مخرجات rowmax/ℓ في FlashAttention: يتضمن FlashAttention مخرجات rowmax/ℓ. نظرًا لأن هذه القيم تمثل الحد الأقصى والإجماليات التراكمية، فإن حجم بياناتها صغير جدًا، مما يؤثر بشكل طفيف على الأداء الكلي لـ FlashAttention. من خلال إعادة استخدام rowmax/ℓ، يمكننا تجنب الحاجة إلى مسار ثانٍ (1-pass)، وبالتالي تعزيز الكفاءة الحسابية الإجمالية.
- إضافة خطوة Online Softmax: نقدم خطوة Online Softmax، التي تبني على FlashAttention من خلال حساب درجات softmax باستخدام الصيغة p=e((kq)−row_max)/ℓ (منطق المسار الواحد). بعد ذلك، من خلال إجراء جمع صفّي (row-wise summation)، يمكننا حساب درجة التشتت S.
في تنفيذ Online Softmax، هناك تفصيلان مهمان يتطلبان عناية خاصة:
- ضبط ترتيب الحلقة: نظرًا لأن S هو مجموع الدرجات بناءً على بُعد k، فمن الضروري ضبط الحساب في Online Softmax من q*k إلى *k**q، مقارنة بآلية الانتباه. يتجنب هذا التعديل الاتصال غير الضروري بين الكتل المختلفة عن طريق تحويل الجمع العمودي إلى جمع صفّي، مما يسمح بإكمال جميع الحسابات داخل السجلات (registers).
- مواقع بداية ونهاية q: نظرًا لوجود الترميز السببي (causal encoding) ومنطق الجمع لأقرب a، ليست هناك حاجة لاجتياز جميع قيم q مع k. من خلال الرجوع إلى مؤشرات القناع المسماة 0/1 في الشكل 7، يمكننا تحديد ما إذا كانت الكتلة المقابلة تتطلب حسابًا بناءً على مواقع بداية ونهاية q، مما يوفر موارد حسابية بشكل كبير. على سبيل المثال، إذا كان طول المطالبة المدخلة 5000 وطول أقرب a هو 256، فإن الحساب التقليدي سيتطلب عمليات 5000**5000. ومع ذلك، من خلال تسجيل مواضع البداية والنهاية، ستكون هناك حاجة فقط إلى عمليات 5000**256، مما يعزز الكفاءة الحسابية.
يؤدي PagedAttention عملاً مشابهًا لعمل FlashAttention، بهدف مشترك هو حساب الدرجة النهائية S. لا يتم تقديم شرح مفصل لـ PagedAttention هنا.
الخاتمة
إصدارنا المعدل، Novita AI Sparse 0.5.1، المبني على vLLM 0.5.1، يدعم حاليًا بشكل أساسي نماذج مثل Llama3–8B و Llama3–70B. يعمل في وضع الاستدلال باستخدام رسومات CUDA ويتم نشره بشكل أساسي على وحدات معالجة رسوميات للمستهلكين مثل RTX 4090. لا يزال هناك مجال كبير للتحسين في الجوانب الهندسية والخوارزمية.
مجرد تنبيه: الميزات التالية غير مدعومة بعد في vLLM-0.6.2:
- Chunked-prefill
- Prefix caching
- FlashInfer وواجهات خلفية أخرى غير FlashAttention
للمهتمين بـ Novita AI Sparse 0.5.1، يمكنكم زيارة Novita AI لتجربتها قبل الآخرين.
المراجع:
[1] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
[2] Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference
[3] SnapKV: LLM Knows What You are Looking for Before Generation
[4] PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling
[5] PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference
[6] MiniCache: KV Cache Compression in Depth Dimension for Large Language Models
[7] Layer-Condensed KV Cache for Efficient Inference of Large Language Models
[8] TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding
[9] CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion
[10] KV-Runahead: Scalable Causal LLM Inference by Parallel Key-Value Cache Generation
