‏شکار حباب‌های خالی GPU در موتور استنتاج Photon

تیم مهندسی Moondream تو این پست توضیح می‌ده که موتور استنتاجشون به اسم Photon چطور استنتاج مدل‌های زبانی-تصویری رو تقریباً بلادرنگ می‌کنه (حدود ۳۳ میلی‌ثانیه روی NVIDIA B200) و تا ۳۵٪ توان عبوری بیشتری موقع تولید توکن به دست میاره. نکته‌ی اصلی نوشته اینه که خیلی وقت‌ها GPU نه به خاطر کمبود کار، بلکه چون CPU هنوز نگفته قدم بعدی چیه، بی‌کار می‌شینه.

مدل‌های زبانی متن رو توکن‌به‌توکن و پشت سر هم تولید می‌کنن؛ هر توکن به توکن قبلی وابسته‌ست، پس این کار ذاتاً ترتیبیه. تو هر دور، GPU حجم زیادی محاسبه‌ی سنگین انجام می‌ده، ولی CPU هم کلی کار جانبی داره: انتخاب اینکه کدوم درخواست بعدی اجرا بشه، آماده کردن متادیتا، برداشتن توکن خروجی و ثبتش. چون کار یک توکن روی GPU کوچیکه ولی این کارهای جانبی هزینه‌ی ثابتی توی هر دور دارن، GPU بخشی از هر حلقه بی‌کار می‌مونه. به گفته‌ی نویسنده به همین وقفه می‌گن حباب GPU.

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

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

مکانیزم دوم «اول اجرا، بعد نمونه‌گیری» ـه. اجرای بعدی به کاری که CPU روی توکن قبلی می‌کنه وابسته نیست، پس زودتر شروع می‌شه؛ ولی توی رمزگشایی محدودشده (که خروجی ساختارمند مثل مختصات یا جعبه می‌ده) ماسکِ توکن‌های مجاز به توکن قبلی وابسته‌ست. برای همین هر تیک زمان‌بند سه فاز داره: اول اجرا رو پرتاب می‌کنن، بعد قدم قبلی رو تثبیت می‌کنن (commit)، و آخر با وضعیت به‌روز، ماسک رو می‌سازن و نمونه‌گیری می‌کنن. به این ترتیب تثبیت از مسیر بحرانی خارج می‌شه.

مکانیزم سوم برخورد با چیزیه که Photon بهش می‌گه «زامبی». برای پرتاب قدم بعدی باید از قبل تصمیم بگیرن کدوم دنباله‌ها تو دسته‌ان، ولی ممکنه یه دنباله همون قدم قبلی به توکن پایان برسه در حالی که تو اجرای بعدی هم هست. به‌جای منطق لغو پیچیده، با دو فیلد ساده (finalized و inflight_refs) کار می‌کنن: دنباله‌ی تموم‌شده نتیجه‌ش صادر می‌شه ولی تا وقتی ارجاع در جریان داره پاک نمی‌شه و تثبیت بعدیش نادیده گرفته می‌شه. به این می‌گن «زودتر نهایی کن، دیرتر آزاد کن».

نویسنده یه مدل هزینه هم می‌ده: با پایپ‌لاین کردن، دوره‌ی هر قدم به سمت اجرا + نمونه‌گیری جمع می‌شه و حباب تقریباً از بین می‌ره. سود این کار هرچی سخت‌افزار سریع‌تر (یا مدل کوچیک‌تر) باشه بیشتره: روی یک 3090 حدود ۱۲٪ ولی روی B200 حدود ۳۵٪ با ۳۲ جریان. تنها هزینه‌ش «مالیات زامبی»ـه که فقط با یک جریان محسوسه (حدود ۱٪) و توی دسته‌ی پر تقریباً محو می‌شه. جمع‌بندی تیم اینه که Photon به خاطر یک تکنیک سریع نیست، بلکه ده‌ها همچین جزئیاتی روی هم انباشته می‌شن.

نکات کلیدی:

• «حباب GPU» یعنی بی‌کار موندن کارت گرافیک تا وقتی CPU کار جانبی‌ش رو تموم کنه
• راه‌حل پایپ‌لاین کردنه: اجرای توکن بعدی همزمان با جمع‌کردن توکن قبلی شروع می‌شه
• سه مکانیزم: اسلات‌های پینگ‌پونگ، «اول اجرا بعد نمونه‌گیری»، و رفع‌وضع زامبی‌ها
• سود روی سخت‌افزار سریع‌تر بیشتره؛ تا ۳۵٪ روی B200 در برابر حدود ۱۲٪ روی 3090
• prefill هم توی همون پایپ‌لاین دو-اسلاته اجرا می‌شه و نیاز به منطق جدا نداره