سقف عملکرد سخت‌افزارهای AMD بسیار بالاتر از آن چیزی است که پیش‌فرض‌های سازنده نشان می‌دهند. اگر از پردازنده MI300X برای مدل‌های ترنسفورمر استفاده می‌کنید، اکنون راهی برای عبور از محدودیت‌های نرم‌افزاری استاندارد AITER v3 پیدا کرده‌اید.
در حالی که AITER v3 استاندارد بهینه‌شده‌ی شرکت AMD است، تیم MoonMath AI با ارائه یک هسته (Kernel) جدید برای توجه (Attention) با دقت bf16، توانسته است از این استاندارد پیشی بگیرد. این انتشار متن‌باز تحت لایسنس MIT ثابت می‌کند که سازمان‌دهی تخصصی حافظه می‌تواند توان عملیاتی (Throughput) بسیار بیشتری از معماری CDNA3 استخراج کند تا آنچه در پیاده‌سازی‌های استاندارد دیده می‌شود. توسعه این پروژه به لطف دسترسی مستقیم به سخت‌افزار (Bare-metal) توسط HotAisle، یکی از ارائه‌دهندگان ابری AMD، میسر شده است.
برای توسعه‌دهندگان، عملیات «توجه» حیاتی‌ترین عملیات ادغام‌شده است که در قلب هر مدل ترنسفورمر قرار دارد و با فرمول ریاضی $\text{softmax}(QK^T/\sqrt{d}) \cdot V$ تعریف می‌شود. در حالی که انویدیا با Triton و cuDNN بر اکوسیستم نرم‌افزاری تسلط دارد، اکوسیستم AMD بر پشته ROCm و HIP متکی است. دستیابی به حداکثر عملکرد در MI300X (مدل gfx942) معمولاً نیازمند دانش عمیق از معماری مجموعه دستورات (ISA) و کنترل دقیق روی هسته‌های واحد پردازش گرافیکی (GPU) است. این هسته خاص به زبان HIP نوشته شده است، نه با اسمبلی دست‌نویس کامل، و منحصراً برای سخت‌افزار gfx942 طراحی شده است.
همان‌طور که در بررسی‌های پیشین ما درباره بهینه‌سازی‌های سطح پایین در مدل‌های زبانی اشاره کردیم، فاصله بین کد سطح بالا و سخت‌افزار اغلب منجر به اتلاف منابع می‌شود. در این پروژه، نویسندگان برای حذف این فاصله، محدودیت‌های سخت‌افزاری را به نقطه قوت تبدیل کرده‌اند.
مفاهیم فنی و محدودیت‌های اجرایی
برای درک عملکرد این هسته، ابتدا باید محدودیت‌های آن را شناخت. یک هسته در واقع برنامه کوچکی است که مستقیماً روی هسته‌های GPU اجرا می‌شود تا محاسباتی خاص را با بیشترین سرعت ممکن انجام دهد. این پیاده‌سازی به‌طور خاص بر توجه پیشرو (Forward Attention) در دقت bf16 تمرکز دارد:
- پشتیبانی از چیدمان: ورودی‌ها را در هر دو حالت BSHD یا BHSD می‌پذیرد که نیاز به جابه‌جایی داده‌ها (Transpose) را به‌طور کامل حذف می‌کند.
- ابعاد: بُعد سر (Head Dimension) روی ۱۲۸ ثابت شده است، اما هر طول توالی، از جمله توجه متقاطع (Cross-Attention)، پشتیبانی می‌شود.
- محدودیت‌ها: نسخه فعلی از ماسک‌های علّی (Causal Masking)، توجه پرس‌وجوی گروهی (GQA) یا دسته‌بندی با طول متغیر (Varlen) پشتیبانی نمی‌کند.
- دقت عددی: کنترل عددی بسیار سخت‌گیرانه است. تمام سه حالت گرد کردن با قوانین AITER مطابقت دارد و هر خروجی محدود در محدوده ۱ bf16 ULP از AITER قرار می‌گیرد. همچنین مدیریت مقادیر NaN و Inf به‌صورت بیت-به-بیت یکسان و تعیین‌پذیر (Deterministic) است.
ترفند اسمبلی «تک‌دستوری»
به نقل از تحلیل فنی marktechpost.com، نوآوری اصلی در استفاده از رپرهای اسمبلی تک‌دستوری است. معمولاً توسعه‌دهندگان با یک انتخاب دشوار روبرو هستند: یا از توابع داخلی کامپایلر برای حفظ نظم کد استفاده کنند (که به قیمت از دست دادن کنترل دقیق تمام می‌شود) یا از اسمبلی خام داخلی (Inline Assembly) استفاده کنند (که به قیمت مدیریت دستی و دشوار ثبات‌ها تمام می‌شود).
تیم MoonMath این مشکل را با بسته‌بندی دقیق یک دستور در یک تابع __device__ __forceinline__ حل کرد. برای مثال، پیاده‌سازی آن‌ها از کد زیر استفاده می‌کند:
__device__ __forceinline__ void asm_mfma(bf16x4_t a, bf16x4_t b, fp32x4_t& c) { asm volatile("v_mfma_f32_16x16x16_bf16 %0, %1, %2, %0" : "+v"(c) : "v"(a), "v"(b)); }
آن‌ها با استفاده از محدودیت‌های اسمبلی پیشرفته مانند +v توانستند ورودی و خروجی انباشتگر (Accumulator) را به یک ثبات عمومی برداری (VGPR) واحد گره بزنند. این کار مانع از تولید دستورات کپی غیرضروری v_mov توسط کامپایلر شده و اجازه می‌دهد تیم، ماشین را دستور به دستور هدایت کند، در حالی که مدیریت تخصیص ثبات‌ها و ردیابی جریان داده همچنان بر عهده کامپایلر است.
بهینه‌سازی‌های معماری
علاوه بر رپرهای اسمبلی، این هسته یک خط لوله اجرایی پیچیده برای واحدهای محاسباتی CDNA3 (که دارای چهار واحد SIMD هستند) پیاده کرده است:
- مدیریت Waveها: به‌جای بلوک چهار-وِیو (four-wave) متداول در کتاب‌های درسی، MoonMath از هشت وِیو در هر بلوک استفاده می‌کند که به دو گروه چهارتایی تقسیم شده‌اند. این گروه‌ها دارای اختلاف فاز هستند؛ یعنی در حالی که یک گروه هسته ماتریسی را اشباع می‌کند، گروه دیگر عملیات سافت‌مکس و بارگذاری داده‌ها را انجام می‌دهد. این استراتژی تضمین می‌کند که هسته ماتریسی هرگز بیکار نماند.
- همگام‌سازی: خط لوله از دو
s_barriersدر هر تکرار استفاده می‌کند؛ یکی در زمان تحویل فاز و دیگری در مرز تکرار، در حالی که انتظارها بر اساس شمارنده (per-counter waits) باقی موارد را مدیریت می‌کنند. این ساختار یادآور تناوب در FlashAttention-3 است، اما به دلیل اینکه جابه‌جایی‌های حافظه در CDNA3 به‌طور پیش‌فرض نامتقارن هستند، نیازی به تقسیم وارپ‌های تولیدکننده/مصرف‌کننده ندارد. - جایگاه حافظه: تیم بر روی حافظه‌های کش خاص سخت‌افزاری اولویت‌بندی کرد. جریان‌های کلیدی (K) از حافظه پهنای‌باند بالا (HBM) به یک حافظه محلی (LDS) با ظرفیت ۳۲ کیلوبایت و با بافر دوگانه (Double-buffered) منتقل می‌شوند که بین هر هشت وِیو مشترک است. مقادیر (V) در کش L1 گرم نگه داشته می‌شوند و پرس‌وجوها (Q) به همراه انباشتگرها در VGPRها قرار می‌گیرند.
- انتخاب MFMA: آن‌ها شکل ۱۶×۱۶×۱۶ را برای ضرب-جمع ماتریسی (MFMA) ترجیح دادند تا شکل ۳۲×۳۲×۸. اگرچه توان عملیاتی هر دو یکسان است، اما تایل کوچک‌تر باعث می‌شود در هر لاین تنها ۴ المان fp32 انباشته شود (در مقابل ۱۶ المان)، که فشار روی VGPRها را کاهش می‌دهد. این امر فضای کافی برای ذخیره تایل سوم Q (3Q) ایجاد می‌کند تا بازاستفاده از داده‌ها و پیش‌خوانی (Prefetching) عمیق‌تر شود.
تحلیل بنچمارک‌ها و نتایج
آزمون‌های انجام شده روی سخت‌افزارهای MI300X نشان می‌دهد که این هسته در تمامی اشکال تست شده و حالت‌های گرد کردن (RTNE، RTNA و RTZ) از AITER v3 پیشی گرفته است. تیم سه حالت گرد کردن را بررسی کرد: RTNE (گرد کردن به نزدیک‌ترین عدد زوج)، RTNA (گرد کردن به نزدیک‌ترین، در صورت تساوی دور از صفر) و RTZ (قطع کردن به سمت صفر).
میانگین هندسی (Geomeans) در کل بررسی‌ها نشان‌دهنده افزایش سرعت ۱.۱۸ برابر در حالت RTNE، ۱.۱۵ برابر در RTNA و ۱.۰۸ برابر در RTZ است. در سناریوهای خاص با بار کاری بالا، این شکاف عملکردی بیشتر می‌شود:
- شکل (۲, ۲۴, ۸۱۹۲, ۱۲۸) در RTNE: زمان اجرا از ۳.۷۹ میلی‌ثانیه (AITER) به ۳.۰۸ میلی‌ثانیه (MoonMath) رسید که معادل ۱.۲۳ برابر افزایش سرعت است.
- شکل (۲, ۲۴, ۱۶۳۸۴, ۱۲۸) در RTNE: زمان اجرا از ۱۴.۶۹ میلی‌ثانیه (AITER) به ۱۱.۶۷ میلی‌ثانیه (MoonMath) کاهش یافت که معادل ۱.۲۶ برابر افزایش سرعت است.
- شکل (۱, ۱۶, ۱۳۱۰۷۲, ۱۲۸) در RTNE: زمان اجرا از ۲۶۹.۲۷ میلی‌ثانیه (AITER) به ۲۳۲.۵۱ میلی‌ثانیه (MoonMath) رسید که معادل ۱.۱۶ برابر افزایش سرعت است.
در مقایسه با Modular MAX، میانگین هندسی این هسته بین ۱.۴۴ تا ۱.۴۹ برابر است و در نقاط اوج به ۱.۵۹ برابر می‌رسد. همچنین تیم یک تقسیم‌بندی دم (Tail KV split) به سبک Flash-Decoding را برای مدیریت دورهای کسری در ۳۰۴ واحد محاسباتی (CU) پردازنده MI300X پیاده کرد. این کار کمک کرد تا شکاف در سخت‌ترین رقابت، یعنی شکل (۴, ۱۶, ۱۶۳۸۴) در حالت RTZ، از ۰.۹۵ برابر به ۱.۰۷ برابر ارتقاء یابد.
کاربرد واقعی: انتشار ویدیو
این نتایج صرفاً بنچمارک‌های مصنوعی نیستند. تیم MoonMath این هسته را در SGLang برای پشتیبانی از LiteAttention در مدل انتشار ویدیو Wan2.1 ادغام کرد. با جایگزینی AITER با پیاده‌سازی liteattention_rocm روی سخت‌افزار MI300X، سرعت تولید سرتاسری (End-to-End) در مدل Wan2.1-T2V-1.3B-Diffusers حدود ۱.۲۳ برابر بهبود یافت، بدون اینکه هیچ کاهش کیفیتی در خروجی‌های بصری گزارش شود.
این نتیجه به‌ویژه برای انتشار ویدیو اهمیت دارد، زیرا چیدمان BSHD استفاده شده در این هسته مستقیماً با تنسورهای مدل‌های انتشار سازگار است و اجازه می‌دهد عملیات توجه متقاطع بدون نیاز به Padding یا Transpose به‌صورت بهینه انجام شود.
این تغییر رویکرد نشان می‌دهد که سقف عملکرد برای سخت‌افزارهای AMD بسیار بالاتر از آن چیزی است که پیش‌فرض‌های سازنده القا می‌کنند. محققان مستقل با کنار گذاشتن توابع داخلی استاندارد و روی آوردن به کنترل دقیق کد عملیاتی (Opcode)، راه‌هایی را برای بهینه‌سازی جریان داده پیدا می‌کنند که کامپایلر اغلب آن‌ها را نادیده می‌گیرد. این امر فشار بر کتابخانه‌های رسمی سازندگان می‌آورد تا در زمینه بازاستفاده از ثبات‌ها و ماندگاری در کش L1 تهاجمی‌تر عمل کنند.
گام بعدی شما
- اگر از محیط‌های PyTorch ROCm استفاده می‌کنید، اکنون می‌توانید این هسته را از طریق pip نصب کنید تا API آن را تست نمایید.
- بررسی کنید که آیا مدل‌های انتشار ویدیو شما می‌توانند از چیدمان BSHD برای حذف عملیات Transpose بهره ببرند یا خیر.
- منتظر به‌روزرسانی‌های بعدی برای پشتیبانی از GQA و ماسک‌های علّی باشید، زیرا این قابلیت‌ها در نسخه فعلی پشتیبانی نمی‌شوند.
اما تأثیر این بهینه‌سازی بر هزینه‌های استنتاج در مقیاس مرکز داده حتی جذاب‌تر است — به تحلیل ما درباره اقتصاد GPUهای نسل جدید مراجعه کنید.
گفتگو