کیمی K3 روی کارتهای AMD؛ ارزونتر از انویدیا
هوش مصنوعیخلاصهٔ کاملتر
به گفتهٔ تیم Wafer، مدلهای متنباز چند ماهه جهش بزرگی داشتن و کیمی K3 نقطهٔ عطف تازهایه؛ ولی این هوش بیشتر با اندازهٔ خیلی بزرگتر هم اومده. K3 با ۲.۸ تریلیون پارامتر بیش از ۱.۵ ترابایت VRAM میخواد، اون هم قبل از جا دادن KV cache برای یک میلیون توکن کانتکست. یعنی حتی یه نود هشتکارته B200 هم جواب نمیده.
تو این مقاله اومده که گزینهها یا نود B300 با ۲۸۸ گیگابایت VRAM روی هر GPU ـه یا دو نود B200 بهصورت TP16. اما MI355X شرکت AMD هم دقیقاً همون ۲۸۸ گیگ حافظه رو داره و بهطور میانگین حدود ۲.۴ برابر ارزونتر از B300 تموم میشه. مشکل همیشگی AMD پشتیبانی نرمافزاریه، ولی چون AMD برای K3 پشتیبانی روز-صفر داده، بیشتر کار از قبل انجام شده بود.
روی بنچمارک ۱۰۲۴ توکن ورودی و ۴۰۰ توکن خروجی، MI355X به ۹۵۲ توکن بر ثانیه در هر نود و ۱۱۸ توکن تکجریانی رسیده؛ بیش از ۳.۸ برابر throughput کل هر نود نسبت به دیپلوی TP16 روی B200. کارت B300 هنوز حدود ۱.۶۵ برابر جلوتره، ولی وقتی قیمت وارد معادله میشه ورق برمیگرده: ۴۸ توکن بر ثانیه به ازای هر دلارِ ساعتی برای MI355X در برابر ۳۳ توکن B300.
بزرگترین اهرم بهبود، speculative decode بود. K3 هیچ draft tensor داخلی نداره، برای همین باید از مدل جانبی Kimi-K3-DSpark استفاده کرد. روی CUDA بیدردسر اجرا میشه، ولی روی ROCm اولین درخواست واقعی، زمانبند sglang رو با این خطا از کار میانداخت:
NameError: name 'top_k_renorm_prob' is not definedنویسنده میگه ریشهش این بود که بیلد ROCm فقط یه کرنل Triton برای top-p رو alias کرده و top_k_renorm_prob تعریفنشده میمونه. راهحل هم یه تابع سادهٔ PyTorch بود نه کرنل سفارشی: مرتبسازی، صفر کردن درایههای اضافه و نرمالسازی دوباره. با همین اصلاح حدود ۲.۲ برابر بهبود تکجریانی و ۱۸٪ افزایش throughput کل به دست اومد.
به نظر نویسنده، decode بیش از اندازه بزرگ شده و چیزی که کاربر واقعاً حسش میکنه TTFT ـه. یه prefill سرد ۱۷۲ هزار توکنی روی MI355X حدود ۵۱ ثانیه طول میکشید در برابر ۲۳ ثانیه روی B300، چون کرنل سریع AITER MLA بهخاطر ناهمخوانی شکل (۱۲ هد توجه در هر رنک بهجای ۴، ۸ یا مضارب ۱۶) اصلاً لود نمیشد. پد کردن ۱۲ به ۱۶ با صفر، prefill رو ۲ تا ۳ برابر سریعتر کرد.
نکات کلیدی:
- کیمی K3 با ۲.۸ تریلیون پارامتر توی یه نود B200 جا نمیشه
- هشت کارت MI355X به ۹۵۲ توکن بر ثانیه در هر نود میرسه
- B300 سریعتره ولی حدود ۲.۴ برابر گرونتر
- رفع باگ top-k renorm روی ROCm فقط چند خط PyTorch بود
- پد کردن هدهای توجه از ۱۲ به ۱۶، prefill رو ۲ تا ۳ برابر سریع کرد
