از CUDA تا MLX: چگونه K-Search دههها تجربه کار با کرنیلها را به اپل سیلیکون میآورد
شکل ۱: نقشه ترجمه بهینهسازی CUDA به MLX. دانش بهینهسازی CUDA میتواند بهجای کپی دستور به دستور، به استراتژیهای بومی معماری MLX ترجمه شود.
ما با عصر جدیدی در محاسبات مواجه هستیم. سختافزارها بهسرعت در حال تغییر هستند — نه تنها پردازندههای گرافیکی سریعتر، بلکه طیف گستردهای از تراشهها از سازندگان مختلف که هرکدام معماری خاص خود را داشته و برای بارهای کاری خاص هوش مصنوعی (AI workloads) تنظیم شدهاند. نرمافزار نیز به همان سرعت تغییر میکند و ابزارهای کدنویسی هوش مصنوعی اکنون در عرض چند دقیقه کدهایی تولید میکنند که چند سال پیش ماهها زمان میبرد.
با تمرکز بخش عمدهای از محاسبات بر هوش مصنوعی، کرنیلهای GPU مولفهای حیاتی برای موفقیت آن هستند. این برنامههای سطح پایین که داخل GPU اجرا میشوند، نیازمند سالها تخصص برای بهینهسازی هستند. انتقال یک کرنیل از سختافزار یک شرکت به شرکت دیگر دشوارتر است و اغلب به معنای بازآفرینی بهینهسازیها از صفر است. برای مثال، اکوسیستم CUDA دههها تجربه ارزشمند را در خود جای داده است (مانند پیادهسازیهای دستی Attention و مدلهای فضای حالت یا SSM). اکوسیستمهای جدیدتر مانند اپل سیلیکون (Apple Silicon) سریع رشد میکنند اما فاقد این عمق از تجربه هستند.
در این پژوهش، ما بررسی میکنیم که آیا این دانش میتواند به صورت خودکار منتقل شود یا خیر. ما چارچوب K-Search (یک چارچوب تکاملی بهینهسازی کرنیل توسعهیافته در برکلی) را با یک بکاند برای MLX — فریمورک یادگیری ماشین اپل برای تراشههای اپل سیلیکون — گسترش دادیم. ما یک لایه ترجمه ساختاریافته CUDA به MLX توسعه دادیم که به K-Search اجازه میدهد از کرنیلهای موجود CUDA به عنوان پایگاه دانش استفاده کرده و آنها را به کرنیلهای باکیفیت برای اپل سیلیکون تبدیل کند.
نتایج ما نشان میدهد که این روش در کرنیل Attention به عملکرد 0.97x نسبت به کرنیل بومی MLX و در کرنیل Mamba SSM به سرعت پیشپرکردن (Prefill) تا ۲۰ برابر سریعتر از پیادهسازی جامعه کاربری (mlx-lm) دست یافته است.
چرا MLX؟
فریمورک MLX اپل از اواخر سال ۲۰۲۳ استقبال گستردهای داشته است. حضور اپل سیلیکون در صدها میلیون مکبوک و مکاستودیو، امکان استنتاج محلی هوش مصنوعی را بدون هزینههای ابری فراهم میکند. معماری حافظه یکپارچه (Unified Memory Architecture) آن را بهویژه برای مدلهای متوسط (۷ تا ۷۰ میلیارد پارامتر) جذاب کرده است.
با این حال شکاف عمدهای وجود دارد: بسیاری از کرنیلهای کلیدی مرتبط با عملکرد مانند paged attention و SSM scan kernels یا وجود ندارند یا بدون تنظیمات اختصاصی عملکرد ضعیفی دارند.
K-Search چیست؟
K-Search یک چارچوب تکاملی برای بهینهسازی کرنیل است. با دریافت یک کرنیل ساده و مشخصات سختافزاری، یک حلقه بهینهسازی تکرارشونده را اجرا میکند: یک مدل زبانی بزرگ (LLM) درباره بهینهسازیها استدلال میکند، مدل تولید کد نمونههای کاندید را میسازد و آن نمونهها روی سختافزار واقعی ارزیابی میشوند.
جستجو توسط یک Spec هدایت میشود که قوانین سختافزاری و الگوهای بهینهسازی را مشخص میکند. حالت استدلال پایدار به صورت یک مدل جهانی (World Model) در قالب یک درخت تصمیمگیری نگهداری میشود.
ساخت بکاند MLX و ترجمه دانش CUDA
برای آوردن K-Search به اپل سیلیکون، ما یک بکاند بومی MLX ایجاد کردیم. لایه ترجمه ما شامل موارد زیر است:
- جدول نگاشت مفاهیم: واژهنامهای از معادلهای CUDA در Metal با محدودیتهای سختافزاری (مانند نگاشت __shared__ به threadgroup memory با حد ۳۲ کیلوبایت).
- راهنماها و الگوهای اختصاصی MLX: کدهای نمونه برای عملیاتی بدون معادل مستقیم مانند ترفند exp2 برای محاسبه سریعتر softmax روی سختافزار اپل.
- ادعاهای قابل بازاستفاده (Assertions): بازتعریف رفتار کرنیلهای حرفهای به عنوان ویژگیهایی که جستجوی تکاملی باید حفظ کند.
ارزیابی عملکرد
ارزیابیها نشان میدهد کرنیل Attention تکاملیافته توانسته است به نزدیک عملکرد کرنیل بومی اپل برسد و در کرنیل Mamba SSM نیز افزایش سرعت ۲۰ برابری در مرحله prefill نسبت به پیادهسازی mlx-lm حاصل شود.