کوهیر با یه مگاکرنل، سرو مدل رو سریعتر از vLLM کرد
هوش مصنوعیخلاصهٔ کاملتر
کوهیر (Cohere) یه سرویسدهندهی جدید برای اجرای مدل زبانیش به اسم North Mini Code معرفی کرده که قلبش یه megakernel هست؛ یعنی بهجای اینکه هر مرحلهی محاسبه (مثل QKV، attention، MoE) رو با یه کرنل جدا اجرا کنن و بین هر کدوم منتظر بمونن، کل یه پاس رو تو یه کرنل واحد و پیوسته اجرا میکنن. روی یه GPU از نوع H100 و با دقت BF16، این روش بین ۱.۲۵ تا ۱.۴۱ برابر از vLLM سریعتره، و کدش رو هم روی گیتهاب گذاشتن.
مشکلی که حل میکنن اینه که decode (مرحلهای که مدل توکنبهتوکن جواب تولید میکنه) بیشتر محدود به پهنای باند حافظهست نه قدرت محاسباتی؛ چون تو هر قدم باید حجم زیادی از وزنهای مدل از HBM (حافظهی اصلی GPU) خونده بشه ولی محاسبهی نسبتا کمی روش انجام میشه. برای North Mini Code با ۳.۳ میلیارد پارامتر فعال، هر قدم decode باید ۶.۶ گیگابایت وزن جابهجا کنه؛ روی H100 با پهنای باند ۳.۳۵ ترابایت بر ثانیه، سقف نظری (Speed-of-Light) حدود ۴۷۰ توکن بر ثانیهست. vLLM فقط ۱۸۵ توکن بر ثانیه یعنی ۳۹٪ این سقف رو میرسه، ولی megakernel کوهیر رو batch size ۱ به ۲۹۲ توکن بر ثانیه (۶۲٪ سقف) میرسه که ۱.۵۸ برابر سریعتر از vLLM هست؛ این فاصله تا کانتکست ۲۵۶ هزار توکنی هم حفظ میشه، بدون افت دقت.
دلیل اصلی کندی روش معمول اینه که بین هر دو کرنل متوالی، همهی SMها (پردازندههای موازی GPU) باید کارشون تموم بشه تا کرنل بعدی شروع بشه. megakernel چهار مشکل رو حل میکنه: هزینهی راهاندازی و همگامسازی بین کرنلها کم میشه؛ wave quantization (وقتی تعداد کارها دقیقا مضربی از تعداد SMها نیست و چندتا SM بیکار میمونن) با پر کردن SMهای بیکار از کار آمادهی عملیات موازی دیگه کاهش پیدا میکنه؛ وابستگیهای الکی بین بخشهای مختلف حذف میشه چون سدهای همگامسازی الان ریزدونهان نه کل-گرید؛ و وزنهای مرحلهی بعدی، حتی قبل از تموم شدن محاسبهی مرحلهی فعلی، شروع به بارگذاری میکنن.
از نظر پیادهسازی، همهی عملیاتها از یه قرارداد واحد پیروی میکنن: هر بلوک ترد ۱۲ warp داره (یه warp یعنی ۳۲ ترد که با هم اجرا میشن) که نقششون کنترلکننده، تولیدکننده، ذخیرهکننده و مصرفکنندهست، و هر تسک با یه توصیفگر ۳۲تایی مشخص میشه که میگه چه کاری باید انجام بشه و منتظر کدوم سیگنال بمونه. وابستگیها هم بهجای مرز کرنل، با شمارندههای ساده تو حافظهی سراسری پیادهسازی شدن:
while (*(volatile const uint32_t*)bar < target) {
__nanosleep(20);
}
__threadfence();
atomicAdd(bar, 1);این مکانیزم ساده باعث میشه صبر کردن و سیگنال دادن هر دو با هزینهی ثابت انجام بشن، فارغ از اینکه چندتا تسک دیگه به این وابستگی وصلن. کوهیر میگه نوشتن megakernel اونقدرا که فکر میکنین سخت نیست و یه دستورالعمل هم برای پورت کردن کرنلهای موجود بهش دادن: از یه GEMM (ضرب ماتریسی) و attention که تنهایی هم سریعن شروع کن، کدتو با قرارداد warp-ها هماهنگ کن، سدهای وابستگی رو اضافه کن، و در نهایت توصیفگر تسکها رو تو یه لیست تو حافظهی سراسری بریز. سختترین بخش، مدیریت درست سدهاست چون یه اشتباه تو تعداد سیگنالها میتونه باعث قفلشدگی (deadlock) بشه که خیلی دیرتر خودشو نشون میده.
زمانبندی تسکها هم عمدتا ایستاست (یه ترتیب از پیش تعیینشده برای هر لایه که بهصورت round-robin بین SMها پخش میشه) بهعلاوهی یه بخش داینامیک برای کارهایی که اندازهشون از قبل معلوم نیست، مثل طول KV هر درخواست تو attention یا اینکه هر توکن به کدوم expert تو MoE میره؛ اونجا هر SM آزاد از یه صف مشترک کار میدزده. آزمایشهای کوهیر نشون داده ترتیب چیدن این مراحل خیلی مهمه: یه چیدمان «درهم» که attention و MoE رو قاطی میکنه فقط ۱ تا ۴٪ کندتر از بهترین ترتیبه، ولی چیدمانی که اول کل attention رو تموم میکنه و بعد میره سراغ MoE، بسته به اندازهی batch بین ۷ تا ۱۹٪ افت سرعت داره.
نکات کلیدی:
- کوهیر یه موتور سرویسدهی برای مدل North Mini Code معرفی کرده که بهجای دهها کرنل جدا، کل پاس رو تو یه megakernel واحد اجرا میکنه.
- روی H100 با BF16، بین ۱.۲۵ تا ۱.۴۱ برابر از vLLM سریعتره؛ رو batch size ۱ به ۲۹۲ توکن بر ثانیه (۶۲٪ سقف نظری پهنای باند) میرسه در برابر ۱۸۵ توکن بر ثانیهی vLLM.
- این سرعت تا کانتکست ۲۵۶ هزار توکنی و بدون افت دقت حفظ میشه؛ کل سیستم از continuous batching و paged attention و یه API سازگار با OpenAI پشتیبانی میکنه.
- طراحی از پروژهی «Look Ma, No Bubbles!» تیم Hazy Research الهام گرفته ولی بدون شبیهسازی صف حافظهی مشترک که به گفتهی کوهیر پیچیده و باگزا بود.
- کد کامل تو یه فایل CUDA نوشته شده و روی گیتهاب در دسترسه.
- ترتیب زمانبندی تسکها روی سرعت خیلی تاثیر داره؛ یه چیدمان بد میتونه تا ۱۹٪ کندتر بشه.
