‏IKBO: چطور متا حافظه و محاسبات بی‌فایده رو از سیستم توصیه‌گر حذف کرد

وقتی کاربر فیدش رو باز می‌کنه، سیستم توصیه‌گر باید صدها تا هزاران آیتم کاندیدا رو امتیازدهی کنه. ورودی‌های مدل به دو دسته تقسیم می‌شن: ویژگی‌های کاربر (مثل تاریخچه جستجو و پروفایل) که برای همه کاندیداهای یه request یکسانه، و ویژگی‌های کاندیدا که برای هر آیتم منحصربه‌فرده. در اصطلاح متا، این دو دسته رو Request-Only یا RO و Non-Request-Only یا NRO می‌نامن.

مشکل اینجاست که لایه‌های تعامل (مثل linear projection و attention) نیاز دارن تنسورها بُعد batch یکسانی داشته باشن. مثلاً وقتی ۱۵ کاربر داریم و ۱۰۲۴ کاندیدا، embedding کاربر باید ~۷۰ بار کپی (broadcast) بشه تا با بُعد batch کاندیداها match کنه. این کپی‌سازی هم حافظه می‌خوره، هم پهنای باند IO رو هدر می‌ده و با افزایش تعداد کاندیداها به‌صورت خطی بدتر می‌شه.

ایده اصلی IKBO اینه که broadcast یه مسئله چیدمان داده‌ست، نه یه ضرورت محاسباتی. به جای اینکه سیستم قبل از kernel تنسورهای تکراری بسازه، هر kernel خودش broadcast رو داخل خودش مدیریت می‌کنه. دو نوع بهینه‌سازی تعریف شده: Type I که از تجزیه ریاضی استفاده می‌کنه تا بخش RO رو مستقل و با batch size کوچک‌تر محاسبه کنه، و Type II که صرفاً جابجایی داده‌های اضافی رو حذف می‌کنه و kernel رو از IO-bound خارج می‌کنه.

برای kernel Linear Compression، متا چهار مرحله بهینه‌سازی پشت سر هم اعمال کرده: تجزیه matmul، هم‌راستاسازی حافظه، ادغام broadcast، و در نهایت fusion چندمرحله‌ای با تخصص‌بندی warp از طریق TLX (Triton Low-Level Extensions). نتیجه نهایی روی H100 SXM5 یه speedup تجمعی ~۴× بوده.

برای kernel Flash Attention، IKBO با حذف تنسورهای broadcast‌شده، kernel رو از IO-bound به compute-bound تبدیل کرده و به ۶۲۱ BF16 TFLOPs رسیده. مقایسه با baseline غیر co-designed یعنی CuTeDSL FA4-Hopper نشون می‌ده که throughput ۲.۴× (فقط kernel) و ۶.۴× (kernel + broadcasting) بهتر شده.

از نظر معماری سیستم، پیاده‌سازی IKBO سه لایه رو لمس می‌کنه: kernel های سفارشی که batch size های نامتقارن RO/NRO رو می‌پذیرن، کامپایلر ML که باید shape range های داینامیک هر operator رو بدونه، و runtime استنتاج که نقشه کاندیدا-به-کاربر رو مستقیم به مدل می‌ده بدون اینکه broadcast رو مادی کنه. این رویکرد end-to-end روی هر دو سخت‌افزار GPU و MTIA در تمام مراحل ranking متا پیاده‌سازی شده.

نکات کلیدی:

• IKBO با مدیریت broadcast داخل kernel، از ساخت تنسورهای تکراری جلوگیری می‌کنه
• دو نوع بهینه‌سازی داره: تجزیه ریاضی (Type I) و حذف جابجایی داده اضافی (Type II)
• kernel Linear Compression روی H100 SXM5 به ~۴× speedup رسیده
• kernel Flash Attention به ۶۲۱ BF16 TFLOPs رسیده و ۶.۴× از baseline سریع‌تره
• تا ۲/۳ کاهش در compute-intensive net latency در مدل‌های co-designed
• روی GPU و MTIA در تمام مراحل ranking متا (از early-stage تا late-stage) استقرار یافته
• ستون اصلی مقیاس‌پذیری مدل Meta Adaptive Ranking Model هست