حافظه KV و PagedAttention: چرا سرور LLM شما کند می‌شود؟

سرور LLM شما کند کار می‌کند.

شما یک مدل 70B را روی چهار پردازنده گرافیکی A100 مستقر کرده‌اید. ساعت ۸ صبح همه چیز خوب به نظر می‌رسد. تا زمان ناهار، تأخیر (latency) دو برابر می‌شود. حافظه خود را بررسی می‌کنید. بیشتر آن توسط "tensor buffers" اشغال شده است. این‌ها در واقع حالت‌های ذخیره‌شده (cached states) از گفتگوهای قدیمی هستند.

این همان مشکل KV cache است. این بزرگترین گلوگاه در سرویس‌دهی LLM در محیط عملیاتی است.

KV cache چیست؟

هر مدل transformer توکن‌ها را یکی پس از دیگری تولید می‌کند. برای ایجاد یک توکن جدید، مدل به تنسورهای Key و Value تمام توکن‌های قبلی نیاز دارد. محاسبه مجدد این‌ها در هر بار بسیار کند است. در عوض، موتور محاسباتی آن‌ها را ذخیره می‌کند. این ذخیره‌سازی همان KV cache است.

مشکل حافظه:

برای یک مدل Llama 3.1 70B، یک توالی ۴۰۹۶ توکنی به تنهایی حدود ۱.۳ گیگابایت حافظه نیاز دارد.

اگر همزمان ۲۵۶ کاربر داشته باشید، به ۳۳۶ گیگابایت حافظه نیاز خواهید داشت. این مقدار بیشتر از ظرفیت چهار پردازنده گرافیکی A100 است. KV cache با چنان سرعتی رشد می‌کند که اغلب بیشتر از خودِ وزن‌های مدل (model weights)، حافظه اشغال می‌کند.

مدیریت حافظه سنتی شکست می‌خورد زیرا:

  • قطعه‌قطعه شدن داخلی (Internal fragmentation): شما فضایی برای ۴۰۹۶ توکن اختصاص می‌دهید اما فقط از ۳۰۰ تای آن استفاده می‌کنید. ۹۳٪ از آن فضا هدر می‌رود.
  • عدم اشتراک‌گذاری: دو کاربر با یک system prompt یکسان، هر کدام نسخه مخصوص به خود را از آن پرامپت ذخیره می‌کنند.
  • حذف تمام یا هیچ (All-or-nothing eviction): وقتی حافظه پر می‌شود، باید کل توالی را به CPU منتقل کنید. این کار باعث توقف (stall) GPU می‌شود.

چگونه PagedAttention این مشکل را حل می‌کند:

PagedAttention مانند یک سیستم‌عامل عمل می‌کند. این روش KV cache را به بلوک‌های کوچک با اندازه ثابت به نام page تقسیم می‌کند.

این کار سه مشکل اصلی را حل می‌کند:

  • تخصیص بر حسب نیاز (On-demand allocation): یک توالی تنها با رشد کردن، pageها را اشغال می‌کند. شما حافظه را برای ظرفیت استفاده نشده هدر نمی‌دهید.
  • پشتیبانی از پیشوند مشترک (Shared prefix support): چندین کاربر می‌توانند از pageهای فیزیکی یکسان برای یک system prompt مشترک استفاده کنند. این کار از منطق "copy-on-write" برای صرفه‌جویی در حجم عظیمی از حافظه استفاده می‌کند.
  • حذف ریزدانه (Fine-grained eviction): وقتی حافظه پر می‌شود، سیستم به جای انتقال توالی‌های عظیم، pageهای کوچک را به CPU منتقل می‌کند.

نتیجه:

استفاده از PagedAttention (تکنولوژی به‌کار رفته در vLLM) می‌تواند نرخ پردازش (throughput) را در مقایسه با روش‌های سنتی ۲ تا ۴ برابر افزایش دهد.

چه زمانی از آن استفاده کنیم:

  • همزمانی بالا (High concurrency).
  • توالی‌هایی با طول‌های متفاوت.
  • پرامپت‌هایی که شروع یکسانی دارند.

چه زمانی از آن صرف‌نظر کنیم:

  • استنتاج محلی (local inference) برای یک کاربر واحد.
  • مدل‌های بسیار کوچک.
  • وظایفی که در آن‌ها طول تمام توالی‌ها دقیقاً یکسان است.

KV Cache و PagedAttention: آن‌ها چه کار می‌کنند و چرا اهمیت دارند؟

در دنیای مدل‌های زبانی بزرگ (LLM)، سرعت و کارایی در زمان استنتاج (Inference) بسیار حیاتی است. اگر شما با چت‌بات‌هایی مثل ChatGPT کار کرده باشید، احتمالاً متوجه شده‌اید که آن‌ها کلمات را یکی پس از دیگری تولید می‌کنند. اما پشت صحنه، مدیریت حافظه برای این فرآیند بسیار پیچیده و چالش‌برانگیز است.

دو مفهوم کلیدی که در این زمینه بسیار مهم هستند، KV Cache و PagedAttention می‌باشند. در این مقاله، بررسی می‌کنیم که این دو چیستند و چرا برای اجرای کارآمد مدل‌های زبانی بزرگ ضروری هستند.

KV Cache چیست؟

برای درک اهمیت KV Cache، ابتدا باید بدانیم که مکانیزم "توجه" (Attention Mechanism) در مدل‌های ترنسفورمر (Transformer) چگونه کار می‌کند.

در هر مرحله از تولید متن، مدل باید به تمام توکن‌های قبلی که تولید شده‌اند نگاه کند تا بتواند توکن بعدی را پیش‌بینی کند. این فرآیند شامل محاسبه مقادیر Key (کلید) و Value (مقدار) برای هر توکن است.

اگر ما در هر مرحله، تمام محاسبات مربوط به توکن‌های قبلی را از ابتدا انجام دهیم، این کار بسیار هزینه‌بر و کند خواهد بود. برای جلوگیری از این محاسبات تکراری، ما مقادیر Key و Value مربوط به توکن‌های قبلی را ذخیره می‌کنیم. به این ذخیره‌سازی، KV Cache می‌گوییم.

با استفاده از KV Cache، مدل در هر مرحله فقط نیاز دارد محاسبات مربوط به توکن جدید را انجام دهد و سپس آن را به مقادیر قبلی که در حافظه ذخیره شده‌اند اضافه کند. این کار سرعت استنتاج را به شدت افزایش می‌دهد.

مشکل کجاست؟ (چالش‌های مدیریت حافظه)

با وجود اینکه KV Cache سرعت را بالا می‌برد، اما یک مشکل بزرگ ایجاد می‌کند: مصرف بسیار بالای حافظه گرافیکی (VRAM).

اندازه KV Cache با طول متن (Context Length) و اندازه مدل رابطه مستقیم دارد. با افزایش طول متن، مقدار حافظه مورد نیاز به صورت بسیار سریع افزایش می‌یابد. این موضوع منجر به دو نوع مشکل اصلی در مدیریت حافظه می‌شود:

۱. قطعه‌قطعه شدن داخلی (Internal Fragmentation)

در روش‌های سنتی، برای هر درخواست (Request)، یک بلوک حافظه بزرگ و پیوسته (Contiguous) به اندازه حداکثر طول متن ممکن (Max Sequence Length) رزرو می‌شود.

مشکل اینجاست که اکثر درخواست‌ها به آن حداکثر طول نیاز ندارند. برای مثال، اگر حافظه برای ۱۰۰۰ توکن رزرو شده باشد اما کاربر فقط ۵۰ توکن تولید کند، ۹۵۰ توکن از فضای رزرو شده بدون استفاده باقی می‌ماند. این فضای هدر رفته، قطعه‌قطعه شدن داخلی نامیده می‌شود.

۲. قطعه‌قطعه شدن خارجی (External Fragmentation)

از آنجایی که حافظه باید به صورت بلوک‌های پیوسته باشد، ممکن است مجموع فضای خالی حافظه برای یک درخواست جدید کافی باشد، اما چون این فضاها در جاهای مختلف حافظه پراکنده هستند و پشت سر هم نیستند، سیستم نمی‌تواند از آن‌ها استفاده کند. این پدیده قطعه‌قطعه شدن خارجی است.

راه حل: PagedAttention

در اینجا PagedAttention وارد عمل می‌شود. این تکنیک که توسط تیم vLLM معرفی شد، از ایده‌ای در سیستم‌عامل‌ها به نام حافظه مجازی (Virtual Memory) و صفحه‌بندی (Paging) الهام گرفته شده است.

در سیستم‌عامل‌ها، حافظه فیزیکی به بلوک‌های کوچکی به نام "صفحه" (Page) تقسیم می‌شود. برنامه‌ها فکر می‌کنند که یک فضای حافظه پیوسته دارند، اما در واقعیت، صفحات آن‌ها می‌تواند در جاهای مختلف حافظه فیزیکی پراکنده باشد.

PagedAttention دقیقاً همین کار را برای KV Cache انجام می‌دهد:

  1. تقسیم به بلوک‌های کوچک: به جای رزرو کردن یک بلوک بزرگ و پیوسته برای هر درخواست، KV Cache به بلوک‌های کوچک و با اندازه ثابت تقسیم می‌شود.
  2. عدم نیاز به پیوستگی: این بلوک‌ها می‌توانند در هر جای خالی از حافظه گرافیکی (VRAM) قرار بگیرند. نیازی نیست که آن‌ها حتماً پشت سر هم باشند.
  3. جدول صفحه (Page Table): یک ساختار داده به نام "جدول صفحه" مدیریت می‌کند که هر توکن در کدام بلوک و کدام قسمت از حافظه ذخیره شده است.

مزایای PagedAttention

  • کاهش چشمگیر هدررفت حافظه: با استفاده از بلوک‌های کوچک، فضای رزرو شده بسیار به میزان واقعی مورد نیاز نزدیک‌تر است، که باعث کاهش قطعه‌قطعه شدن داخلی می‌شود.
  • افزایش توان عملیاتی (Throughput): چون حافظه بسیار بهینه‌تر مدیریت می‌شود، می‌توان درخواست‌های بسیار بیشتری را به طور همزمان در یک GPU اجرا کرد.
  • اشتراک‌گذاری حافظه: PagedAttention اجازه می‌دهد که بخش‌هایی از KV Cache (مانند Promptهای مشترک) بین درخواست‌های مختلف به اشتراک گذاشته شوند، که این امر مصرف حافظه را باز هم کمتر می‌کند.

نتیجه‌گیری

KV Cache برای سرعت استنتاج ضروری است، اما مدیریت ناکارآمد آن می‌تواند باعث هدررفت عظیم منابع شود. PagedAttention با الهام از مفاهیم سیستم‌عامل، این مشکل را حل کرده و اجازه می‌دهد مدل‌های زبانی بزرگ با کارایی بسیار بالاتر و استفاده بهینه‌تر از سخت‌افزار اجرا شوند. این تکنولوژی سنگ بنای بسیاری از سیستم‌های استنتاج مدرن و مقیاس‌پذیر است.


Source: https://dev.to/tech_nuggets/kv-cache-and-pagedattention-what-they-do-and-why-they-matter-jce

Optional learning community: https://t.me/GyaanSetuAi