بهینه‌سازی استنتاج مدل‌های زبانی: از کش KV تا PagedAttention

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

استنتاج در مدل‌های زبانی به دو مرحله‌ی کاملاً متفاوت تقسیم می‌شود: فاز پیش‌تکمیل (Prefill) که محدود به توان پردازشی است، و فاز دیکود (Decode) که محدود به پهنای باند حافظه است. برای حل چالش‌های این دو فاز، تکنیک‌هایی نظیر کش کلید-مقدار (KV Cache)، دسته‌بندی پیوسته (Continuous Batching)، پجد اتنشن (PagedAttention)، کش پیشوند و دیکودینگ حدسی توسعه یافته‌اند. این روش‌ها مانع از بیکار ماندن هسته‌های پردازشی و تکه‌تکه‌شدن حافظه می‌شوند.

با توجه به اینکه تأمین GPU قدرتمند یا اجاره‌ی آن از ارائه‌دهندگان خارجی برای کاربران داخلی بسیار پرهزینه و دشوار است، بهینه‌سازی حداکثری هر گیگابایت حافظه و زمان‌بندی دقیق استنتاج اهمیتی دوچندان پیدا می‌کند. در این درس با معماری موتورهای مدرنی همچون vLLM و SGLang آشنا می‌شوید و یاد می‌گیرید چگونه با همان سخت‌افزار موجود، نرخ پردازش را چندین برابر کنید.

استنتاج مدل‌های زبانی بزرگ (LLM) در دو فاز خلاصه می‌شود: پیش‌تکمیل (Prefill) پرامپت شما را به‌صورت موازی پردازش می‌کند و گلوگاه آن توان پردازشی (Compute-bound) است. دیکود (Decode) توکن‌ها را یکی‌یکی تولید می‌کند و گلوگاه آن پهنای باند حافظه (Memory-bound) است. تمام تکنیک‌های بهینه‌سازی دقیقاً روی یک یا هر دوی این فازها تمرکز دارند.

اهداف یادگیری

• پیاده‌سازی کش کلید-مقدار (KV cache) برای حذف محاسبات تکراری در طول فرآیند تولید خودبازگشتی توکن‌ها
• تشریح تفاوت فازهای Prefill و Decode در استنتاج مدل‌های زبانی و دلیل تفاوت گلوگاه‌های آن‌ها (محدود به پردازش در برابر محدود به حافظه)
• پیاده‌سازی مفاهیم دسته‌بندی پیوسته (Continuous Batching) و پجد اتنشن (PagedAttention) جهت بیشینه‌سازی بهره‌وری GPU زیر بار درخواست‌های هم‌زمان
• مقایسه‌ی تکنیک‌های بهینه‌سازی استنتاج (کش کلید-مقدار، دیکودینگ حدسی و فلش اتنشن) و تحلیل مصالحه‌ها بین نرخ پردازش و تاخیر

صورت مسئله

فرض کنید مدل Llama 3 70B را روی 4 کارت گرافیک A100 مستقر کرده‌اید. یک کاربر تکی سرعتی در حدود 50 توکن بر ثانیه دریافت می‌کند؛ عددی بسیار خوب و سریع. ناگهان 100 کاربر به‌صورت هم‌زمان به این اندپوینت درخواست می‌فرستند. در این لحظه نرخ پردازش برای هر کاربر تا 3 توکن بر ثانیه سقوط می‌کند. شما ماهیانه 25,000 دلار برای هزینه‌ی GPU پرداخت می‌کنید، در حالی که سیستم پاسخی کندتر از سرعت تایپ یک انسان تحویل می‌دهد.

بین شرایط 1 کاربر و 100 کاربر، خود مدل هیچ تغییری نکرده است: وزن‌ها همان است، معماری یکسان است و محاسبات ریاضی فرقی نکرده است. چیزی که تغییر کرده، نحوه‌ی زمان‌بندی کارها است. استنتاج خام و ساده‌انگارانه بیش از 90% از ظرفیت پردازشی GPU را هدر می‌دهد. کاربری که منتظر دریافت توکن شماره‌ی 47 است، یک اسلات کامل از بچ را اشغال نگه می‌دارد، در حالی که گذرگاه حافظه‌ی GPU بین محاسبات ضرب ماتریسی بیکار مانده است. در همین فاصله، پرامپت 2,000 توکنیِ یک کاربر جدید می‌توانست این زمان‌های تلف‌شده را با محاسبات مفید پر کند.

این یک مشکل مقیاس‌پذیری نیست؛ یک مشکل زمان‌بندی است. تکنیک‌هایی که در این درس بررسی می‌کنیم — کش کلید-مقدار، دسته‌بندی پیوسته، PagedAttention، دیکودینگ حدسی و کش پیشوند (Prefix Caching) — تفاوت میان یک قبض ماهیانه‌ی 25,000 دلاری استنتاج با یک قبض 5,000 دلاری برای مدیریت همان میزان ترافیک است.

اجرای Llama 3 70B با فریم‌ورک vLLM روی 4 کارت A100-80GB در هم‌زمانی پایین به حدود 50 توکن بر ثانیه به ازای هر کاربر می‌رسد، و با تکیه بر دسته‌بندی پیوسته و PagedAttention، در بار 100 کاربر هم‌زمان نیز سرعتی پایدار بین 15 تا 25 توکن بر ثانیه به ازای هر کاربر فراهم می‌کند. بدون این بهینه‌سازی‌ها، همان سخت‌افزار در چنین باری سرعتی در حدود 5 توکن بر ثانیه برای هر کاربر ارائه می‌دهد. سخت‌افزار یکسان، مدل یکسان، اما 4 برابر نرخ پردازش بالاتر.

مفاهیم پایه

پیش‌تکمیل در برابر دیکود

هر درخواست استنتاج در مدل‌های زبانی از دو فاز کاملاً مجزا عبور می‌کند:

فاز پیش‌تکمیل (Prefill) کل پرامپت ورودی را پردازش می‌کند. در این مرحله تمام توکن‌ها مشخص هستند، بنابراین اتنشن را می‌توان روی کل دنباله به‌صورت موازی محاسبه کرد. این عملیات در عمل یک ضرب ماتریسی بزرگ است و هسته‌های GPU را کاملاً درگیر نگه می‌دارد. گلوگاه این فاز، توان پردازشی است؛ یعنی سخت‌افزار شما در هر ثانیه چند عملیات ممیز شناور (FLOPS) می‌تواند انجام دهد. یک کارت A100 توان 312 ترافلاپس (با دقت BF16) دارد. اجرای فاز پیش‌تکمیل برای یک پرامپت 4,096 توکنی روی مدل 70B حدود 400 میلی‌ثانیه روی یک A100 تکی زمان می‌برد.

فاز دیکود (Decode) توکن‌های خروجی را تک‌به‌تک تولید می‌کند. هر توکن جدید با تمام توکن‌های قبلی اتنشن برقرار می‌کند، اما در هر گذر رو به جلو (forward pass) فقط یک توکن ساخته می‌شود. ماتریس‌های وزن دقیقاً همان اندازه‌ی فاز پیش‌تکمیل را دارند، اما این بار آن‌ها را به جای ماتریس، در یک بردار تکی ضرب می‌کنید. هسته‌های GPU کار محاسباتی خود را ظرف چند میکروثانیه تمام می‌کنند و سپس منتظر می‌مانند تا دسته‌ی بعدی وزن‌ها از حافظه فراخوانی شوند. گلوگاه در اینجا پهنای باند حافظه است؛ یعنی با چه سرعتی می‌توانید وزن‌های مدل را از حافظه‌ی HBM به واحدهای پردازشی منتقل کنید. یک کارت A100 پهنای باندی معادل 2 ترابایت بر ثانیه دارد. حجم یک مدل 70B با دقت FP16 برابر 140 گیگابایت است. یک دور خواندن کامل این مدل از حافظه 70 میلی‌ثانیه طول می‌کشد؛ این یعنی حداقل زمان لازم برای تنها یک گام دیکود.

نسبت عملیات به بایت (ops:byte ratio) که شدت محاسباتی (Arithmetic Intensity) هم نامیده می‌شود، این توازن را توصیف می‌کند. این معیار نشان می‌دهد به ازای هر بایتی که از حافظه بارگذاری می‌شود، چه تعداد عملیات محاسباتی انجام می‌دهید.

ops:byte ratio = FLOPs per token / bytes read from memory

در فاز پیش‌تکمیل با دسته‌ای متشکل از 4,096 توکن، شما به ازای هر وزنِ بارگذاری‌شده حدود 4,096 عملیات ضرب و جمع تجمعی (Multiply-Accumulate) انجام می‌دهید. این نسبت بالا است، بنابراین سیستم محدود به پردازش است. در فاز دیکود با اندازه‌ی بچ 1، به ازای هر وزنِ خوانده‌شده تقریباً 1 عملیات انجام می‌دهید. این نسبت پایین است، بنابراین سیستم محدود به حافظه می‌شود.

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

کش کلید-مقدار (KV Cache)

در لایه‌ی اتنشن، بردار پرس‌وجو (Query) هر توکن روی بردار کلید (Key) و مقدار (Value) تمام توکن‌های قبلی اعمال می‌شود. بدون کش کردن، برای تولید توکن N-ام مجبورید محاسبات تصویر کردن (Projection) کلید و مقدار را برای تمام N-1 توکنِ قبل دوباره تکرار کنید. با این روش، توکن اول هنگام تولید توکن دوم تصویر می‌شود، دوباره برای توکن سوم، و بار دیگر برای توکن چهارم. وقتی به توکن 1,000 برسید، توکن اول را مجموعاً 999 بار از نو محاسبه و تصویر کرده‌اید.

کش کلید-مقدار، بردارهای کلید و مقدار محاسبه‌شده برای تمام توکن‌های گذشته را ذخیره می‌کند. هنگام تولید توکن N-ام، شما فقط کلید و مقدار همان توکن N را حساب می‌کنید و آن را به K و V کش‌شده از توکن‌های 1 تا N-1 می‌چسبانید.

فرمول محاسبه‌ی حافظه‌ی مورد نیاز کش کلید-مقدار:

KV cache size = 2 * num_layers * num_kv_heads * head_dim * seq_len * bytes_per_param

برای مدل Llama 3 70B (شامل 80 لایه، 8 سر کلید-مقدار با معماری GQA، بعد سر معادل 128 و دقت BF16):

per token: 2 * 80 * 8 * 128 * 2 bytes = 327,680 bytes = 320 KB
at 4,096 tokens: 320 KB * 4,096 = 1.28 GB
at 128K tokens: 320 KB * 131,072 = 40 GB

یک گفت‌وگوی ساده با طول زمینه‌ی 128K در مدل Llama 3 70B به‌تنهایی 40 گیگابایت حافظه‌ی کش کلید-مقدار مصرف می‌کند؛ یعنی نیمی از کل حافظه‌ی یک کارت A100. اگر 100 کاربر هم‌زمان داشته باشید که هر کدام زمینه‌ای با 4K توکن داشته باشند، صرفاً نگهداری این کش نیازمند 128 گیگابایت حافظه است. به همین دلیل، مدیریت کش کلید-مقدار مهم‌ترین چالش در بهینه‌سازی استنتاج است.

دسته‌بندی پیوسته (Continuous Batching)

بچینگ ایستا (Static batching) صبر می‌کند تا دسته‌ای شامل N درخواست برسد، همه‌ی آن‌ها را با هم پردازش می‌کند و تا زمانی که همه‌ی آن‌ها تمام نشوند، هیچ درخواست جدیدی نمی‌پذیرد. اگر کاربری به 500 توکن و کاربری دیگر فقط به 10 توکن نیاز داشته باشد، درخواست کوتاه‌تر پس از پایان کارش، به مدت 490 گام دیکود بیکار می‌ماند.

دسته‌بندی پیوسته (که دسته‌بندی در سطح تکرار یا Iteration-level batching نیز نامیده می‌شود)، به محض اتمام هر کدام از درخواست‌ها، درخواست جدیدی را وارد بچ می‌کند. وضعیت بچ در هر گام دیکود دوباره ارزیابی می‌شود. درخواستی که پس از 10 توکن کارش تمام شده، بلافاصله جایش را به درخواستی می‌دهد که در صف انتظار بوده است.

میزان بهبود نرخ پردازش به تفاوت طول خروجی درخواست‌ها وابسته است. اگر طول پاسخ‌ها یکسان باشد، دسته‌بندی پیوسته عملکردی مشابه بچینگ ایستا دارد. اما در حالت تفاوت طول خروجی‌ها (که حالت رایج در کاربردهای واقعی است)، دسته‌بندی پیوسته می‌تواند نرخ پردازش را بین 2 تا 5 برابر افزایش دهد، زیرا هیچ‌یک از اسلات‌های GPU خالی نمی‌مانند.

پجد اتنشن (PagedAttention)

کش کلید-مقدار برای هر درخواست یک بلوک پیوسته در حافظه (Contiguous) است. با ورود و خروج درخواست‌ها، حافظه دچار تکه‌تکه‌شدگی (Fragmentation) می‌شود؛ دقیقاً مانند مسئله‌ی خرد شدن حافظه‌ی رم در سیستم‌های عامل. درخواستی با طول 4K توکن به 1.28 گیگابایت حافظه‌ی پیوسته نیاز دارد. حتی اگر در مجموع 2 گیگابایت فضای خالی داشته باشید، ممکن است 1.28 گیگابایت فضای پیوسته در دسترستان نباشد. در چنین حالتی یا حافظه هدر می‌رود، یا ناچارید درخواست را رد کنید.

تکنیک PagedAttention (طراحی‌شده در فریم‌ورک vLLM) ایده‌ی حافظه‌ی مجازی سیستم‌های عامل را به کش کلید-مقدار می‌آورد. به جای تخصیص یک بلوک پیوسته برای هر درخواست، حافظه در قالب «صفحه‌ها» یا پیج‌هایی با اندازه‌ی ثابت (معمولاً 16 توکن) تخصیص داده می‌شود. این پیج‌ها می‌توانند در هر نقطه‌ای از حافظه‌ی فیزیکی GPU باشند. یک جدول پیج (Page Table) موقعیت منطقی توکن‌های هر درخواست را به موقعیت فیزیکی پیج‌ها نگاشت می‌کند.

همچنین PagedAttention امکان کپی در هنگام نوشتن (Copy-on-write) را برای پیشوندهای مشترک فراهم می‌کند. اگر 50 درخواست از یک پرامپت سیستم یکسان استفاده کنند، پیج‌های کش کلید-مقدار مربوط به آن پرامپت فقط یک بار ذخیره شده و توسط هر 50 درخواست ارجاع داده می‌شوند. فقط زمانی که مسیر درخواست‌ها از هم جدا شود (پیام‌های متفاوت کاربران)، پیج‌های اختصاصی برای آن‌ها تخصیص می‌یابد. این سازوکار مصرف حافظه را در برنامه‌هایی با پرامپت سیستم مشترک به‌شدت کاهش می‌دهد.

طبق آمار vLLM، به لطف PagedAttention هدررفت حافظه تقریباً به صفر می‌رسد (حدود 4% در برابر 60 تا 80 درصد در روش‌های سنتی تخصیص پیوسته).

دیکودینگ حدسی (Speculative Decoding)

دیکود به این دلیل کند است که باید به‌صورت متوالی انجام شود؛ یعنی یک توکن می‌سازید، آن را بازخورد می‌دهید و توکن بعدی را می‌سازید. اما اگر بتوانید 5 توکن بعدی را با هزینه‌ای ناچیز حدس بزنید و سپس صحت همه‌ی آن‌ها را یک‌باره ارزیابی کنید چه می‌شود؟

دیکودینگ حدسی از یک مدل پیش‌نویس (Draft Model) کوچک و سریع برای تولید K توکن پیشنهادی استفاده می‌کند. سپس مدل هدف (Target Model) اصلی و بزرگ، تمام K توکن کاندید را در یک گذر رو به جلوی تکی پردازش می‌کند (فرآیندی که درست مانند پیش‌تکمیل، به‌صورت موازی، محدود به پردازش و بهینه اجرا می‌شود). اگر مدل هدف پیش‌بینی‌های مدل پیش‌نویس را تأیید کند، شما تمام K توکن را تنها در زمان یک گذر مدل هدف به دست آورده‌اید. اگر پیش‌بینی در موقعیت j نادرست باشد، توکن‌های 1 تا j-1 را می‌پذیرید و بقیه را دور می‌ریزید.

میزان افزایش سرعت به نرخ پذیرش (Acceptance Rate) وابسته است؛ یعنی پیش‌بینی‌های مدل پیش‌نویس چند وقت یک‌بار با مدل هدف همخوانی دارند. اگر مدل Llama 3 8B نقش پیش‌نویس را برای Llama 3 70B بازی کند، نرخ پذیرش معمولاً روی متون طبیعی بین 70 تا 85 درصد است که به 2 تا 3 برابر سرعت بیشتر در فاز دیکود ترجمه می‌شود.

سه رویکرد اصلی در دیکودینگ حدسی وجود دارد:

روش EAGLE یک سر خودبازگشتی کوچک روی وضعیت‌های پنهان (Hidden states) مدل هدف آموزش می‌دهد. این روش با استفاده از ویژگی‌های لایه‌ی یکی‌مانده‌به‌آخرِ مدل هدف، امبدینگ توکن بعدی را پیش‌بینی می‌کند. از آنجا که این تکنیک مستقیماً روی بازنمایی‌های خود مدل هدف کار می‌کند (نه یک مدل خارجی)، با کمترین سربار حافظه به نرخ پذیرش بالاتری می‌رسد. نسخه‌ی EAGLE-2 یک درخت کاندیدای پویا نیز اضافه کرده که تعداد توکن‌های پیش‌نویس را بر اساس بافت متن تنظیم می‌کند.

دیکودینگ حدسی مبتنی بر N-gram جدولی از ادامه‌های n-gram برگرفته از زمینه‌ی فعلی یا یک پیکره‌ی متنیِ آماده را نگه می‌دارد. اگر پیش‌نویس با آنچه پیش‌تر در همان گفت‌وگو ظاهر شده مطابقت داشته باشد (الگوهای تکراری، کد، خروجی ساختاریافته)، حدس‌ها بدون هیچ‌گونه سربار شبکه‌ی عصبی اجرا می‌شوند. میانگین نرخ پذیرش این روش پایین‌تر است، اما هزینه‌ی هر بار حدس زدن عملاً رایگان تمام می‌شود.

دیکودینگ حدسی از نظر ریاضی کاملاً دقیق است؛ یعنی توزیع احتمالاتی خروجی، تطابق کامل و صددرصدی با توزیع مدل هدف دارد. این روش تقریب نیست، زیرا مرحله‌ی اعتبارسنجی تضمین می‌کند که هر توکنِ پذیرفته‌شده دقیقاً با همان احتمالی ثبت شود که مدل هدف تعیین می‌کرد.

کش پیشوند (Prefix Caching)

بسیاری از درخواست‌ها پیشوند مشترکی دارند؛ پرامپت سیستم یک چت‌بات، زمینه‌ی بازیابی‌شده در RAG یا مجموعه‌ای از مثال‌های چندشات (Few-shot). بدون استفاده از کش پیشوند، سیستم مجبور است کش کلید-مقدار را برای این توکن‌های مشترک در تک‌تک درخواست‌ها از نو محاسبه کند.

کش پیشوند این مقادیر را برای پیشوندهای متداول ذخیره کرده و میان درخواست‌ها بازاستفاده می‌کند. با رسیدن یک درخواست جدید با پیشوند شناخته‌شده، سرور ورودی‌های کش‌شده را ارجاع می‌دهد (یا کپی می‌کند) و تنها کلید و مقدار پسوند یکتا را محاسبه می‌کند.

اگر یک پرامپت سیستم 2,000 توکنی بین تمام درخواست‌ها مشترک باشد، کش پیشوند حدود 400 میلی‌ثانیه زمان پیش‌تکمیل را برای هر درخواست حذف می‌کند. در مقیاس 100 درخواست بر ثانیه، این یعنی ذخیره‌ی 40 ثانیه از پردازش GPU در هر ثانیه؛ مقداری که بیشتر از ظرفیت پردازشی کامل یک GPU مجزا است.

رویکرد RadixAttention در فریم‌ورک SGLang، کش پیشوند را به کمک درخت پیشوندی رادیکس (Radix Tree یا Trie) پیاده‌سازی می‌کند که پیشوندها را بر اساس محتوای توکن‌هایشان ایندکس می‌کند. هر درخواستی که با یکی از شاخه‌های درخت همخوانی داشته باشد، کش کلید-مقدار را رایگان دریافت می‌کند. ساختار درختی حتی امکان تطبیق جزئی را هم می‌دهد؛ اگر 1,500 توکن از 2,000 توکن پرامپت شما با ورودی کش مطابقت داشته باشد، آن 1,500 توکن بازاستفاده شده و سیستم فقط 500 توکن باقی‌مانده را پردازش می‌کند.

موتورهای استنتاج (Inference Engines)

در محیط‌های عملیاتی، سه موتور بر بازار سرورهای استنتاج مدل‌های زبانی مسلط هستند:

vLLM گزینه‌ی استاندارد و نقطه‌ی شروع پیش‌فرض است. این ابزار از گسترده‌ترین طیف مدل‌ها پشتیبانی می‌کند، روی انواع کارت‌های گرافیک (NVIDIA، AMD، Intel) اجرا می‌شود و با استفاده از PagedAttention در کنار دسته‌بندی پیوسته، نرخ پردازش بسیار بالایی دارد. ارائه‌ی API سازگار با OpenAI به شما اجازه می‌دهد آن را مستقیماً جایگزین هرگونه فراخوانی API قبلی کنید.

SGLang بر پایه‌ی همان مبانی vLLM توسعه یافته اما RadixAttention را برای کش پیشوند و یک زبان اختصاصی دامنه برای برنامه‌های ساختاریافته‌ی مدل زبانی اضافه کرده است. اگر بار کاری شما بر پایه‌ی گفت‌وگوهای چند نوبتی (Multi-turn)، استفاده از ابزارها یا دیکود مقید (Constrained Decoding مانند خروجی JSON یا تولید بر اساس عبارات باقاعده) باشد، SGLang به لطف استفاده‌ی مجدد از پیشوندها معمولاً بین 2 تا 5 برابر سریع‌تر از vLLM عمل می‌کند.

TensorRT-LLM مدل‌ها را به کرنل‌های شدیداً بهینه‌شده‌ی کارت‌های NVIDIA کامپایل می‌کند. این فریم‌ورک عملیات مختلف را ادغام می‌کند (برای مثال اتنشن، لایه‌های خطی و توابع فعال‌ساز در قالب یک کرنل واحد اجرا می‌شوند)، روی نسل H100 از دقت FP8 بهره می‌برد و برای محیط‌های عملیاتی با NVIDIA Triton Inference Server ادغام می‌شود. این فریم‌ورک بالاترین نرخ پردازش تک‌کارت را روی سخت‌افزارهای NVIDIA ارائه می‌دهد، اما راه‌اندازی آن پیچیده‌تر است و فقط روی کارت‌های NVIDIA کار می‌کند.

آمارهای دنیای واقعی برای Llama 3 70B (روی 4 کارت A100-80GB با فرمت BF16):

چارچوب نسبت عملیات به بایت (Ops:Byte)

چیزی را که اندازه نگیرید، نمی‌توانید بهینه‌سازی کنید. نسبت ops:byte مشخص می‌کند که آیا درگیر محدودیت پردازش هستید یا محدودیت پهنای باند حافظه؛ و همین نسبت معلوم می‌کند کدام بهینه‌سازی‌ها اهمیت دارند.

Compute roof: peak FLOPS of the GPU
Memory roof:  peak bandwidth * ops:byte ratio

وقتی نسبت ops:byte پایین است (مانند فاز دیکود با بچ‌های کوچک)، سیستم با سقف پهنای باند حافظه برخورد می‌کند. افزودن توان پردازشی خام (فرکانس بالاتر یا هسته‌های بیشتر) کمکی نمی‌کند. در این حالت باید حجم داده‌های خوانده‌شده از حافظه را کم کنید (مثلاً با کوانتایزیشن یا فشرده‌سازی کش کلید-مقدار) یا اندازه‌ی بچ را بالا ببرید تا هزینه‌ی خواندن وزن‌ها روی کارهای مفید بیشتری سرشکن شود.

وقتی نسبت ops:byte بالا است (مانند فاز پیش‌تکمیل یا بچ‌های بسیار بزرگ)، سیستم با سقف توان پردازشی مواجه می‌شود. در اینجا بهبود پهنای باند حافظه دردی را دوا نمی‌کند؛ بلکه به پردازنده‌های گرافیکی سریع‌تر، فیوژن کرنل‌ها یا کاهش دقت محاسباتی نیاز دارید تا ترافلاپس بیشتری استخراج کنید.

نقطه‌ی تلاقی و تعادل روی کارت A100 تقریباً در نسبت 156 رخ می‌دهد (312 ترافلاپس تقسیم بر 2 ترابایت بر ثانیه). در مقادیر زیر 156، کار شما محدود به حافظه است و در مقادیر بالاتر از 156، محدود به پردازش خواهید بود. دسته‌بندی پیوسته با قرار دادن توکن‌های بیشتر در هر تکرار، فاز دیکود را به سمت این نقطه‌ی تعادل هل می‌دهد.

پیاده‌سازی از صفر

گام 1: پیاده‌سازی کش کلید-مقدار از صفر

در اینجا یک کش کلید-مقدار چندسر می‌سازیم که بردارهای کلید و مقدار را به ازای هر لایه و هر سر ذخیره می‌کند و الگوی رشد حافظه را نشان می‌دهد.

import numpy as np

class KVCache:
    def __init__(self, num_layers, num_heads, head_dim, max_seq_len, dtype=np.float16):
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        self.dtype = dtype

        self.k_cache = np.zeros(
            (num_layers, num_heads, max_seq_len, head_dim), dtype=dtype
        )
        self.v_cache = np.zeros(
            (num_layers, num_heads, max_seq_len, head_dim), dtype=dtype
        )
        self.seq_len = 0

    def update(self, layer_idx, new_keys, new_values):
        num_new = new_keys.shape[1]
        end = self.seq_len + num_new
        self.k_cache[layer_idx, :, self.seq_len:end, :] = new_keys
        self.v_cache[layer_idx, :, self.seq_len:end, :] = new_values
        return (
            self.k_cache[layer_idx, :, :end, :],
            self.v_cache[layer_idx, :, :end, :]
        )

    def advance(self, num_tokens):
        self.seq_len += num_tokens

    def memory_bytes(self):
        return self.k_cache.nbytes + self.v_cache.nbytes

    def used_bytes(self):
        per_token = 2 * self.num_layers * self.num_heads * self.head_dim * np.dtype(self.dtype).itemsize
        return per_token * self.seq_len

گام 2: پیاده‌سازی اتنشن با استفاده از کش کلید-مقدار

پیاده‌سازی یک لایه‌ی اتنشن چندسر (Multi-Head Attention) ساده‌شده که از کش کلید-مقدار برای مراحل دیکود استفاده می‌کند.

def scaled_dot_product_attention(query, keys, values):
    head_dim = query.shape[-1]
    scores = np.matmul(query, keys.transpose(0, 1, 3, 2)) / np.sqrt(head_dim)
    seq_len_q = scores.shape[-2]
    seq_len_k = scores.shape[-1]
    if seq_len_q > 1:
        mask = np.triu(np.ones((seq_len_q, seq_len_k), dtype=np.float32), k=seq_len_k - seq_len_q + 1)
        scores = scores + mask * (-1e9)
    max_scores = np.max(scores, axis=-1, keepdims=True)
    exp_scores = np.exp(scores - max_scores)
    attn_weights = exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)
    return np.matmul(attn_weights, values)


class MultiHeadAttention:
    def __init__(self, d_model, num_heads):
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        scale = np.sqrt(2.0 / d_model)
        self.W_q = np.random.randn(d_model, d_model).astype(np.float32) * scale
        self.W_k = np.random.randn(d_model, d_model).astype(np.float32) * scale
        self.W_v = np.random.randn(d_model, d_model).astype(np.float32) * scale
        self.W_o = np.random.randn(d_model, d_model).astype(np.float32) * scale

    def forward(self, x, kv_cache=None, layer_idx=0):
        batch, seq_len, d_model = x.shape
        Q = np.matmul(x, self.W_q).reshape(batch, seq_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
        K = np.matmul(x, self.W_k).reshape(batch, seq_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
        V = np.matmul(x, self.W_v).reshape(batch, seq_len, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)

        if kv_cache is not None:
            K_full, V_full = kv_cache.update(layer_idx, K[0], V[0])
            K = K_full[np.newaxis, :, :, :]
            V = V_full[np.newaxis, :, :, :]
            if seq_len == 1:
                kv_cache.advance(1)

        attn_out = scaled_dot_product_attention(Q, K, V)
        attn_out = attn_out.transpose(0, 2, 1, 3).reshape(batch, -1, d_model)
        return np.matmul(attn_out, self.W_o)

گام 3: شبیه‌ساز دسته‌بندی پیوسته

این کد تفاوت رفتار زمان‌بندی را میان بچینگ ایستا و دسته‌بندی پیوسته شبیه‌سازی می‌کند.

import heapq

class Request:
    def __init__(self, request_id, prompt_tokens, output_tokens, arrival_step):
        self.request_id = request_id
        self.prompt_tokens = prompt_tokens
        self.output_tokens = output_tokens
        self.arrival_step = arrival_step
        self.tokens_generated = 0
        self.start_step = None
        self.end_step = None

    def is_done(self):
        return self.tokens_generated >= self.output_tokens


def simulate_static_batching(requests, batch_size):
    step = 0
    completed = []
    queue = list(requests)
    queue.sort(key=lambda r: r.arrival_step)

    while queue:
        batch = []
        while queue and len(batch) < batch_size:
            r = queue.pop(0)
            r.start_step = max(step, r.arrival_step)
            batch.append(r)

        if batch:
            step = max(step, max(r.start_step for r in batch))
            max_output = max(r.output_tokens for r in batch)
            for r in batch:
                r.tokens_generated = r.output_tokens
                r.end_step = step + max_output
            step += max_output
            completed.extend(batch)

    return completed


def simulate_continuous_batching(requests, batch_size):
    step = 0
    completed = []
    queue = sorted(requests, key=lambda r: r.arrival_step)
    queue_idx = 0
    active = []
    waiting = []

    while queue_idx < len(queue) or active or waiting:
        while queue_idx < len(queue) and queue[queue_idx].arrival_step <= step:
            waiting.append(queue[queue_idx])
            queue_idx += 1

        while waiting and len(active) < batch_size:
            r = waiting.pop(0)
            r.start_step = step
            active.append(r)

        if not active:
            if waiting:
                step += 1
                continue
            elif queue_idx < len(queue):
                step = queue[queue_idx].arrival_step
                continue
            else:
                break

        for r in active:
            r.tokens_generated += 1

        done = [r for r in active if r.is_done()]
        for r in done:
            r.end_step = step + 1
            completed.append(r)
        active = [r for r in active if not r.is_done()]

        step += 1

    return completed


def batching_stats(completed):
    latencies = [r.end_step - r.arrival_step for r in completed]
    total_time = max(r.end_step for r in completed) - min(r.arrival_step for r in completed)
    total_tokens = sum(r.output_tokens for r in completed)
    return {
        "avg_latency": np.mean(latencies),
        "p50_latency": np.median(latencies),
        "p99_latency": np.percentile(latencies, 99),
        "total_time": total_time,
        "throughput": total_tokens / total_time if total_time > 0 else 0,
    }

گام 4: پیاده‌سازی کش پیشوند

پیاده‌سازی یک کش پیشوند مبتنی بر ساختار درخت پیشوندی (Trie) که مقادیر کش کلید-مقدار را برای پیشوندهای مشترک ذخیره می‌کند.

class TrieNode:
    def __init__(self):
        self.children = {}
        self.kv_data = None
        self.hit_count = 0


class PrefixCache:
    def __init__(self, max_entries=1000):
        self.root = TrieNode()
        self.max_entries = max_entries
        self.total_entries = 0
        self.hits = 0
        self.misses = 0

    def _walk(self, token_ids):
        node = self.root
        depth = 0
        for tid in token_ids:
            if tid not in node.children:
                break
            node = node.children[tid]
            depth += 1
        return node, depth

    def lookup(self, token_ids):
        node, depth = self._walk(token_ids)
        if depth > 0:
            self.hits += 1
            current = self.root
            for tid in token_ids[:depth]:
                current = current.children[tid]
                current.hit_count += 1
            kv_entries = []
            current = self.root
            for tid in token_ids[:depth]:
                current = current.children[tid]
                if current.kv_data is not None:
                    kv_entries.append(current.kv_data)
            return depth, kv_entries
        self.misses += 1
        return 0, []

    def insert(self, token_ids, kv_per_token):
        node = self.root
        for i, tid in enumerate(token_ids):
            if tid not in node.children:
                if self.total_entries >= self.max_entries:
                    return i
                node.children[tid] = TrieNode()
                self.total_entries += 1
            node = node.children[tid]
            if i < len(kv_per_token):
                node.kv_data = kv_per_token[i]
        return len(token_ids)

    def hit_rate(self):
        total = self.hits + self.misses
        return self.hits / total if total > 0 else 0.0

گام 5: شبیه‌ساز دیکودینگ حدسی

شبیه‌سازی فرآیند دیکودینگ حدسی بر پایه‌ی مدل پیش‌نویس و هدف با قابلیت تنظیم نرخ پذیرش توکن‌ها.

class DraftModel:
    def __init__(self, vocab_size, acceptance_rate=0.8):
        self.vocab_size = vocab_size
        self.acceptance_rate = acceptance_rate

    def generate(self, context, num_tokens):
        tokens = np.random.randint(0, self.vocab_size, size=num_tokens)
        return tokens

    def get_probs(self, context, token):
        probs = np.random.dirichlet(np.ones(self.vocab_size))
        return probs


class TargetModel:
    def __init__(self, vocab_size):
        self.vocab_size = vocab_size

    def get_probs(self, context, tokens=None):
        if tokens is not None:
            return [np.random.dirichlet(np.ones(self.vocab_size)) for _ in tokens]
        return np.random.dirichlet(np.ones(self.vocab_size))


def speculative_decode(draft_model, target_model, context, num_speculative=5,
                       draft_cost=1.0, target_cost=10.0, verify_cost=12.0):
    total_tokens = 0
    total_cost = 0.0
    accepted_counts = []
    context = list(context)

    max_tokens = 100

    while total_tokens < max_tokens:
        draft_tokens = draft_model.generate(context, num_speculative)
        total_cost += draft_cost * num_speculative

        target_probs = target_model.get_probs(context, draft_tokens)
        total_cost += verify_cost

        accepted = 0
        for i, token in enumerate(draft_tokens):
            draft_p = draft_model.get_probs(context + list(draft_tokens[:i]), token)
            target_p = target_probs[i]

            r = np.random.random()
            acceptance_prob = min(1.0, target_p[token] / (draft_p[token] + 1e-10))

            if r < draft_model.acceptance_rate:
                accepted += 1
                context.append(token)
                total_tokens += 1
            else:
                new_token = np.random.choice(draft_model.vocab_size, p=target_p)
                context.append(new_token)
                total_tokens += 1
                break

        accepted_counts.append(accepted)

        if accepted == num_speculative:
            bonus_probs = target_model.get_probs(context)
            bonus_token = np.random.choice(draft_model.vocab_size, p=bonus_probs)
            context.append(bonus_token)
            total_tokens += 1

    sequential_cost = total_tokens * target_cost
    return {
        "total_tokens": total_tokens,
        "speculative_cost": total_cost,
        "sequential_cost": sequential_cost,
        "speedup": sequential_cost / total_cost if total_cost > 0 else 1.0,
        "avg_accepted": np.mean(accepted_counts),
        "acceptance_rate": np.mean(accepted_counts) / num_speculative,
    }


def compare_speculation_strategies(vocab_size=1000, num_trials=20):
    results = {}

    for name, acceptance_rate, spec_tokens in [
        ("Draft-target (8B->70B)", 0.78, 5),
        ("EAGLE", 0.85, 6),
        ("N-gram", 0.50, 4),
        ("No speculation", 0.0, 0),
    ]:
        if spec_tokens == 0:
            results[name] = {
                "speedup": 1.0,
                "acceptance_rate": 0.0,
                "avg_accepted": 0.0,
            }
            continue

        trial_results = []
        for _ in range(num_trials):
            draft = DraftModel(vocab_size, acceptance_rate=acceptance_rate)
            target = TargetModel(vocab_size)
            context = list(np.random.randint(0, vocab_size, size=10))
            result = speculative_decode(draft, target, context, num_speculative=spec_tokens)
            trial_results.append(result)

        results[name] = {
            "speedup": np.mean([r["speedup"] for r in trial_results]),
            "acceptance_rate": np.mean([r["acceptance_rate"] for r in trial_results]),
            "avg_accepted": np.mean([r["avg_accepted"] for r in trial_results]),
        }

    return results

گام 6: پروفایلر حافظه‌ی کش کلید-مقدار

محاسبه‌ی دقیق میزان حافظه‌ی مصرفی کش کلید-مقدار بر اساس پیکربندی مدل‌های واقعی.

MODEL_CONFIGS = {
    "Llama-3-8B": {
        "num_layers": 32, "num_kv_heads": 8, "head_dim": 128,
        "model_params_b": 8, "gqa": True,
    },
    "Llama-3-70B": {
        "num_layers": 80, "num_kv_heads": 8, "head_dim": 128,
        "model_params_b": 70, "gqa": True,
    },
    "Llama-3-405B": {
        "num_layers": 126, "num_kv_heads": 8, "head_dim": 128,
        "model_params_b": 405, "gqa": True,
    },
    "Mistral-7B": {
        "num_layers": 32, "num_kv_heads": 8, "head_dim": 128,
        "model_params_b": 7, "gqa": True,
    },
    "GPT-4-est": {
        "num_layers": 120, "num_kv_heads": 96, "head_dim": 128,
        "model_params_b": 1800, "gqa": False,
    },
}


def kv_cache_memory(config, seq_len, dtype_bytes=2):
    per_token = 2 * config["num_layers"] * config["num_kv_heads"] * config["head_dim"] * dtype_bytes
    total = per_token * seq_len
    return {
        "per_token_bytes": per_token,
        "per_token_kb": per_token / 1024,
        "total_bytes": total,
        "total_mb": total / (1024 ** 2),
        "total_gb": total / (1024 ** 3),
    }


def memory_budget(config, gpu_memory_gb, model_dtype_bytes=2, kv_dtype_bytes=2):
    model_memory_gb = config["model_params_b"] * 1e9 * model_dtype_bytes / (1024 ** 3)
    overhead_gb = gpu_memory_gb * 0.1
    available_for_kv = gpu_memory_gb - model_memory_gb - overhead_gb

    if available_for_kv <= 0:
        return {"error": "Model does not fit in GPU memory", "model_memory_gb": model_memory_gb}

    per_token = 2 * config["num_layers"] * config["num_kv_heads"] * config["head_dim"] * kv_dtype_bytes
    max_tokens = int(available_for_kv * (1024 ** 3) / per_token)

    return {
        "gpu_memory_gb": gpu_memory_gb,
        "model_memory_gb": round(model_memory_gb, 1),
        "overhead_gb": round(overhead_gb, 1),
        "available_for_kv_gb": round(available_for_kv, 1),
        "max_total_tokens": max_tokens,
        "max_users_at_2k": max_tokens // 2048,
        "max_users_at_4k": max_tokens // 4096,
        "max_users_at_32k": max_tokens // 32768,
    }

استفاده از ابزارهای واقعی

استفاده به همراه vLLM:

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Meta-Llama-3-70B-Instruct",
    tensor_parallel_size=4,
    enable_prefix_caching=True,
    max_model_len=8192,
    gpu_memory_utilization=0.9,
)

params = SamplingParams(temperature=0.7, max_tokens=256)
outputs = llm.generate(["Explain inference optimization in one paragraph."], params)

استفاده از SGLang برای بهره‌گیری از کش پیشوند و خروجی‌های ساخت‌یافته:

import sglang as sgl

@sgl.function
def classify(s, text):
    s += sgl.system("You are a classifier. Output JSON only.")
    s += sgl.user(f"Classify this text: {text}")
    s += sgl.assistant(sgl.gen("result", regex=r'\{"label": "(positive|negative|neutral)"\}'))

runtime = sgl.Runtime(model_path="meta-llama/Meta-Llama-3-70B-Instruct", tp_size=4)
sgl.set_default_backend(runtime)

results = classify.run_batch([
    {"text": "This product is amazing!"},
    {"text": "Terrible experience."},
    {"text": "It was okay I guess."},
])

استفاده به همراه TensorRT-LLM:

import tensorrt_llm
from tensorrt_llm.runtime import ModelRunner

runner = ModelRunner.from_dir("./llama-70b-trt-engine/", rank=0)

outputs = runner.generate(
    batch_input_ids=[tokenizer.encode("Explain KV caching.")],
    max_new_tokens=256,
    temperature=0.7,
)

خروجی این درس

این درس فایل زیر را تولید می‌کند:

• outputs/skill-inference-optimization.md -- یک مهارت برای عیب‌یابی و بهینه‌سازی سرورهای استنتاج مدل‌های زبانی

تمرین‌ها

1. اسکریپت پروفایلر کش کلید-مقدار را تغییر دهید تا اثر کوانتایزیشن کش به فرمت‌های FP16، FP8 و INT4 را با هم مقایسه کند. برای مدل Llama 3 70B با طول زمینه‌ی 4K، حداکثر تعداد کاربران هم‌زمان را برای هر حالت روی 4 کارت A100-80GB محاسبه کنید. کوانتایزیشن کش به فرمت INT4 باید ظرفیت تعداد کاربران را تقریباً 4 برابر کند.

2. شبیه‌ساز دسته‌بندی پیوسته را توسعه دهید تا میزان بهره‌وری GPU (نسبت اسلات‌های پرشده‌ی بچ در هر گام) را ثبت کند. نمودار این بهره‌وری را در طول زمان برای هر دو حالت بچینگ ایستا و پیوسته با 50 درخواست که طول خروجی آن‌ها از توزیع پارتو پیروی می‌کند (shape=1.5, scale=20)، رسم کنید. دسته‌بندی پیوسته باید نرخ بهره‌وری را بالای 80% نگه دارد.

3. ساختار کش کلید-مقدار را بر اساس اتنشن پرس‌وجوی گروه‌بندی‌شده (Grouped-Query Attention یا GQA) پیاده کنید؛ حالتی که در آن num_kv_heads < num_query_heads است. مدل Llama 3 70B دارای 64 سر پرس‌وجو ولی تنها 8 سر کلید-مقدار است. میزان صرفه‌جویی در حافظه را نسبت به اتنشن چندسر استاندارد محاسبه کنید (کاهش 8 برابری حجم کش کلید-مقدار).

4. یک سیستم کش پیشوند بسازید که از سیاست آزادسازی LRU (کمترین استفاده‌ی اخیر) استفاده کند. مقدار max_entries را روی 500 بگذارید و 1,000 درخواست تولید کنید که 60% آن‌ها یکی از 5 پیشوند پرکاربرد را به اشتراک دارند. نرخ هیت کش (Hit Rate) را ثبت و با حالت کش نامحدود مقایسه کنید. با یک الگوریتم تخلیه‌ی مناسب، نرخ هیت باید بالای 55% باقی بماند.

5. شبیه‌ساز دیکودینگ حدسی را گسترش دهید تا بر اساس ساختار درختی (مشابه رویکرد EAGLE-2) کار کند. به جای یک دنباله‌ی خطی از K توکن پیش‌نویس، درختی از گزینه‌ها بسازید (برای مثال 2 شاخه در هر یک از 3 سطح = 8 گزینه‌ی برگ). مجموع توکن‌های پذیرفته‌شده در هر دور اعتبارسنجی را نسبت به روش خطی مقایسه کنید.

اصطلاحات کلیدی

برای مطالعه‌ی بیشتر

• Kwon et al., "Efficient Memory Management for Large Language Model Serving with PagedAttention" (2023) -- مقاله‌ی اصلی تیم vLLM که مدیریت پیج‌بندی‌شده‌ی کش کلید-مقدار را معرفی کرد و اکنون استاندارد اصلی سرویس‌دهی استنتاج است.
• Leviathan et al., "Fast Inference from Transformers via Speculative Decoding" (2023) -- مقاله‌ی پایه‌ای که اثبات کرد روش حدس و ارزیابی، بدون تغییر در توزیع آماری مدل هدف، سرعت دیکود را 2 تا 3 برابر افزایش می‌دهد.
• Li et al., "EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty" (2024) -- روشی که با آموزش یک لایه روی ویژگی‌های خود مدل هدف به جای استفاده از مدل مجزا، نرخ پذیرش حدس‌ها را ارتقا می‌دهد.
• Zheng et al., "SGLang: Efficient Execution of Structured Language Model Programs" (2024) -- معرفی ساختار RadixAttention برای کش پیشوند و مدلی برنامه‌نویسی برای اجرای کارآمد برنامه‌های چندفراخوانی مدل زبانی.
• Williams et al., "Roofline: An Insightful Visual Performance Model for Multicore Architectures" (2009) -- مقاله‌ی اصلی مدل سقف (Roofline) که چارچوب نسبت عملیات به بایت را برای تحلیل گلوگاه‌های پردازشی و حافظه تدوین کرد.