PagedAttention

PagedAttention خوارزمية للانتباه ومدير للذاكرة يخدم نماذج اللغة الكبيرة (LLM)، ويخزّن ذاكرة KV المؤقتة لكل طلب في كتل صغيرة ثابتة الحجم توضع في أي موضع ضمن ذاكرة GPU، ويتعقبها عبر جدول كتل خاص بكل طلب، على غرار طريقة نظام التشغيل في تعيين صفحات الذاكرة الافتراضية. قدّم ووسوك كوون وجوهان لي وزملاؤهما في جامعة كاليفورنيا في بيركلي وجامعات شريكة هذه التقنية إلى جانب محرك خدمة vLLM ضمن ورقة إدارة فعالة للذاكرة لخدمة نماذج اللغة الكبيرة باستخدام PagedAttention، التي عُرضت في SOSP 2023. ويتيح التخلص من التجزئة الناجمة عن تخصيص الذاكرة المؤقتة المتجاورة لوحدة GPU واحدة الاحتفاظ بعدد أكبر بكثير من التسلسلات في الوقت نفسه، وتفيد الورقة بأن الإنتاجية أعلى بمقدار 2-4 مرات مقارنةً بـ FasterTransformer وOrca عند كمون مماثل.

مشكلة الذاكرة التي تعالجها#

أثناء توليد النصوص، يحتفظ المحوّل بمتجهات المفاتيح والقيم لكل رمز سابق كي لا يعيد حسابها. وتنمو هذه الذاكرة المؤقتة بمقدار مدخلة واحدة لكل رمز في كل طبقة، وقد يشغل تسلسل واحد من نموذج LLaMA ذي 13 مليار معامل ما يصل إلى 1.7 GB. ويكون طول المخرجات مجهولًا عند وصول الطلب، لذا كانت أنظمة الخدمة السابقة تحجز لكل طلب منطقة متجاورة بحجم يكفي لأقصى طول ممكن. وكان معظم تلك المنطقة يبقى فارغًا، كما كانت الفجوات بين المناطق أصغر من أن يعاد استخدامها. وقد قاس فريق vLLM أن هذه الأنظمة تهدر 60-80% من ذاكرة KV المؤقتة، ما كان يحد من حجم الدفعة قبل نفاد القدرة الحاسوبية بوقت طويل.

آلية عمل تقسيم الذاكرة إلى صفحات#

يقسّم PagedAttention الذاكرة المؤقتة إلى كتل، تحتوي كل منها على مفاتيح وقيم عدد ثابت من الرموز، 16 افتراضيًا في vLLM. ولا يتلقى الطلب كتلة جديدة إلا عندما تمتلئ كتلته الأخيرة، وتُخصّص له من أي موضع تتوافر فيه كتلة حرة، بينما يسجّل جدول كتل الطلب ترتيبها. وتتبع نواة آلية الانتباه ذلك الجدول لجمع المفاتيح والقيم كتلةً تلو الأخرى، لذا لا يلزم أن تكون الذاكرة المؤقتة للتسلسل متجاورة. ويقتصر الهدر على الجزء غير الممتلئ في نهاية الكتلة الأخيرة لكل تسلسل، أي أقل من 4% وفق قياسات vLLM.

يخزّن هذا المثال التوضيحي في PyTorch مفاتيح طلب واحد في كتل متناثرة من مجموعة مشتركة، ويقرأها عبر جدول الكتل الخاص بالطلب:

import torch

block_size, head_dim = 16, 64
key_pool = torch.zeros(64, block_size, head_dim)  # physical blocks shared by every request
free_blocks = [41, 7, 19, 2, 33]  # free blocks are scattered across the pool

# A 40-token request needs ceil(40 / 16) = 3 blocks, wherever they happen to be free
keys = torch.randn(40, head_dim)
num_blocks = -(-len(keys) // block_size)
block_table = [free_blocks.pop(0) for _ in range(num_blocks)]
for pos, key in enumerate(keys):
    key_pool[block_table[pos // block_size], pos % block_size] = key

# The attention kernel reads the request's keys back through its block table
gathered = key_pool[block_table].reshape(-1, head_dim)[: len(keys)]
print(f"block table {block_table}, keys intact: {torch.equal(gathered, keys)}")
print(f"unused slots: {num_blocks * block_size - len(keys)} paged vs {2048 - len(keys)} with a 2,048-token reservation")

يطبع جدول الكتل [41, 7, 19] وkeys intact: True، ثم 8 خانات غير مستخدمة للطلب المقسّم إلى صفحات، مقابل 2,008 خانات لمنطقة حُجزت لـ 2,048 رمزًا.

مشاركة الكتل بين التسلسلات#

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

يفيد مؤلفو vLLM بأن هذه المشاركة تقلل استخدام الذاكرة في العينات المتوازية والبحث الحزمي بنسبة تصل إلى 55%، ما يتيح زيادة الإنتاجية حتى 2.2 مرة. وتمتد الفكرة إلى الطلبات المختلفة عبر تجزئة الكتل كاملةً، وهي الطريقة التي ينفّذ بها vLLM تخزين المطالبات مؤقتًا.

تقنيات ذات صلة#

يحدد PagedAttention مواضع تخزين المفاتيح والقيم المخزّنة مؤقتًا، ويُدمج مع أساليب تغيّر كمية البيانات المطلوب تخزينها أو سرعة قراءتها:

• الانتباه السريع: يقسّم حساب الانتباه إلى مربعات بحيث لا تُكتب مصفوفة الدرجات كاملةً مطلقًا في الذاكرة الرئيسية لـ GPU. وتشغّل المحركات الحالية نوى بأسلوب FlashAttention تقرأ من الكتل المقسّمة إلى صفحات.
• التجميع المستمر للدفعات: يقبل طلبًا جديدًا فور انتهاء طلب آخر، ويعيد تقسيم الصفحات كتل التسلسل المنتهي إلى المجموعة الحرة على الفور، ما يتيح مساحة للطلب الجديد.
• انتباه الاستعلامات المجمّعة: يشارك رؤوس المفاتيح والقيم بين رؤوس الاستعلامات، ما يقلّص محتوى كل كتلة.
• تكميم الذاكرة المؤقتة FP8: يخزّن المفاتيح والقيم بدقة 8 بت، ما يقارب خفض ذاكرة التخزين المؤقت إلى النصف مقارنةً بالتنسيقات ذات 16 بت.

تمرّر واجهة Ultralytics لنماذج اللغة الكبيرة وسيطات الطلب دون تغيير، لذا فإن استدعاء إكمالات المحادثة باستخدام n=3 عبر base_url إلى خادم vLLM محلي يعيد ثلاث عينات تتشارك الكتل المخزّنة مؤقتًا للمطالبة، ولا تستهلك ذاكرة جديدة إلا تتمّاتها التي تتباعد عن بعضها.