کرنل‌های Triton سرعت استنتاج مدل‌های کم‌رتبه را ۳.۷۶ برابر کرد

تصور کنید مدلی با دقت بالا داشته باشید که به‌دلیل محدودیت حافظه، روی سخت‌افزارهای لبه (Edge) به‌شدت کند اجرا می‌شود. اکنون با استفاده از کرنل‌های سفارشی، این سد حافظه‌ای شکسته شده و سرعت استنتاج در دستگاه‌های کوچک تا ۳.۷۶ برابر افزایش یافته است. این پیشرفت در مقاله‌ای در «سی‌وپن‌مین symposium بین‌المللی محاسبات موازی و توزیع‌شده با کارایی بالا» (Proceedings of the 35th International Symposium on High-Performance Parallel and Distributed Computing) منتشر شده و شکاف میان تئوری‌های فشرده‌سازی و سرعت واقعی سیستم را پر می‌کند.

مشکل اصلی اینجاست که مدل‌های ترنسفورمر (Transformer) حافظه بسیار زیادی می‌طلبند. برای مثال، مدل Llama-70B برای ذخیره وزن‌های نیم‌دقت (half-precision) به بیش از ۱۴۰ گیگابایت حافظه نیاز دارد که بسیار بیشتر از ظرفیت ۸۰ گیگابایتی اکثر پردازنده‌های گرافیکی تجاری پرچم‌دار است. این فاصله باعث شده است که محققان به سراغ روش‌های فشرده‌سازی کم‌رتبه (Low-Rank) و کم‌رتبهٔ بلوکی (BLR) بروند تا بتوانند این مدل‌های غول‌پیکر را روی سخت‌افزارهای کوچک‌تر و ارزان‌تر جای دهند.

در حالی که تکنیک‌هایی مانند Monarch و BLAST دقت بالاتری نسبت به روش SVD سنتی دارند (چون ساختارهای وزنی را موثرتر جذب می‌کنند)، اما اغلب در ارائه سرعت واقعی در دنیای حقیقی شکست می‌خورند. طبق مستندات این پژوهش، این روش‌ها به‌طور متناقضی در طول استنتاج چند-توکنی (multi-token inference) — که مرحله حیاتی پیش‌پُرکردن (pre-fill) در مدل‌های زبانی بزرگ است — باعث کندی می‌شوند. دلیل این امر جابجایی عظیم داده‌های میانی است که پهنای باند محدود واحد پردازش گرافیکی (GPU) را اشباع و خفه می‌کند.

دیوار حافظه: چرا فشرده‌سازی باعث کندی می‌شود؟

در دستگاه‌هایی با حافظهٔ کش L2 کوچک، مانند NVIDIA A40 (۶ مگابایت) یا Jetson Orin Nano (۲ مگابایت)، تانسورهای میانی لایه‌های BLR به‌قدری بزرگ هستند که نمی‌توانند روی تراشه (on-chip) باقی بمانند. برای مثال، در یک لایه تصویرسازی (projection) در مدل Llama-7B، یک لایه BLAST تانسوری میانی با اندازه تقریبی ۳۲ مگابایت (با دقت bf16 و ابعاد ۱۶ × ۱۰۲۴ × ۱۰۲۴) ایجاد می‌کند. چون این حجم در کش L2 نمی‌گنجد، GPU مجبور است به‌طور مکرر داده‌ها را از DRAM که بسیار کندتر است فراخوانی کند.

این وضعیت باعث تغییر در «مدل سقف» (Roofline Model) می‌شود. در اینجا سیستم از حالت «محدود به محاسبات» (Compute-bound) که توسط عملیات خام FLOPs محدود شده، به حالت «محدود به حافظه» (Memory-bound) تغییر وضعیت می‌دهد که در آن پهنای باند گلوگاه اصلی است. پژوهش‌ها نشان داد در سناریوهای چند-توکنی، در حالی که روش‌های متراکم و کم‌رتبه سنتی همچنان Compute-bound باقی می‌مانند، روش‌های Monarch و BLAST به زیر نقطه شکست سقف (Roofline breakpoint) سقوط می‌کنند و این منجر به جهش‌های شدید در تاخیر (Latency Spikes) می‌شود.

سازوکار دقیق ساختارهای (B)LR

برای درک عمیق این گلوگاه، باید ساختارهای وزنی خاص به‌کار رفته را بررسی کرد:

• وزن‌های متراکم (Dense): یک ماتریس $W \in \mathbb{R}^{i\times o}$ با $i \times o$ پارامتر که نیازمند $n \times i \times o$ عملیات FLOP است. این روش بیشترین قدرت بیان را دارد اما از نظر حافظه بسیار سنگین است.
• کم‌رتبه (LR): وزن‌ها به صورت $W = VU$ تجزیه می‌شوند که پارامترها را به $r(i + o)$ و محاسبات را به $nr(i + o)$ FLOP کاهش می‌دهد. با این حال، در نرخ‌های فشرده‌سازی بالا، دقت مدل به‌شدت افت می‌کند.
• Monarch: وزن‌های متراکم را به $b_2 \times b_1$ بلوک با رتبه یکسان ($r'$) تقسیم می‌کند. این ساختار از $b_1b_2r'(p + q)$ پارامتر و $nb_1b_2r'(p + q)$ FLOP استفاده می‌کند و معمولاً تحت بودجه FLOP یکسان، قوانین مقیاس‌پذیری (Scaling Laws) بهتری نسبت به شبکه‌های متراکم دارد.
• BLAST: یک ساختار BLR تعمیم‌یافته است که در آن بلوک‌ها ماتریس‌های $V_l$ و $U_k$ را به اشتراک می‌گذارند اما ماتریس‌های قطری $S_{l,k}$ منحصر‌به‌فردی دارند. این روش به $r(p + q + b_1b_2)$ پارامتر و $nr(p + q + b_1b_2)$ FLOP نیاز دارد و معمولاً در یک فاکتور فشرده‌سازی (CF) یکسان، بالاترین صحت را ارائه می‌دهد.

مهندسی راهکار با Triton

پژوهشگران برای حل این مشکل، مجموعه‌ای از بهینه‌سازی‌های سخت‌افزار-آگاه را با استفاده از زبان Triton شرکت OpenAI پیاده کردند. آن‌ها ابتدا «ادغام کامل» (Full Fusion) را آزمایش کردند، اما دریافتند که این روش فقط برای رتبه‌های بسیار کوچک (کمتر یا برابر با ۱۲۸) که مربوط به نرخ‌های فشرده‌سازی شدید (بیشتر از ۸ برابر برای لایه‌های Llama-7B) است، کارآمد است. برای رتبه $r = 256$، ادغام کامل در تمام ابعاد ویژگی، به‌دلیل اینکه تایلینگ تک‌بعدی (1-D tiling) موازی‌سازی را محدود می‌کند، به‌طور مداوم کندتر از مدل‌های متراکم بود.

برای لایه‌های Monarch، تیم سه بهینه‌سازی کلیدی را اجرا کرد:
۱. تغییر چیدمان وزن‌ها (Weight Re-layout) ❶: آن‌ها چیدمان حافظه تانسور وزن استاتیک $\mathcal{V}$ (ذخیره شده به صورت $b_1 \times (r'b_2) \times p$) را به گونه‌ای تغییر دادند که بعد میانی ابتدا در امتداد $r'$ و سپس $b_2$ پیوسته باشد. این کار باعث حذف جایگشت (permutation) غیرضروری $r' \leftrightarrow b_2$ شد.
۲. ادغام جایگشت (Permutation Fusion) ❷: جایگشت $b_2 \leftrightarrow b_1$ مستقیماً در اولین کرنل ضرب ماتریسی دسته‌ای ($\text{bmm}$) ادغام شد. این کار از طریق محاسبه شاخص $b_2$ و تنظیم آفست $r'$ در یک تک کرنل Triton محقق شد.
۳. حذف جایگشت نهایی (Final Permutation Avoidance) ❸: با پیش-جایگشت سطر‌های وزن‌های استاتیک بعدی به‌صورت آفلاین، نیاز به جایگشت خروجی نهایی (تبدیل $(b_2, n, q)$ به $(n, q, b_2)$) در زمان اجرا کاملاً حذف شد.

در لایه‌های BLAST، تیم دو استراتژی متمایز را آزمایش کرد:

• ادغام جزئی bmm (Partial bmm Fusion) ❹: این روش جایگشت میانی بین $\mathcal{V}$ و $\mathcal{S}$ را حذف کرد و مانع از تجسم خروجی اولین $\text{bmm}$ در حافظه سراسری (global memory) شد. با این حال، این کار مستلزم نگاشت $\text{bmm}$ دوم به هسته‌های CUDA به عنوان یک ضرب خارجی دسته‌ای بود که باعث از دست رفتن توان عملیاتی ۱۶ برابری هسته‌های تنسور (Tensor Cores) می‌شد.
• ادغام صرفاً جایگشتی با بهینه‌سازی هسته‌های تنسور (Permutation-only Fusion) ❺: این استراتژی برنده بود. با ترانهاده کردن خروجی $\text{dot}()$ در داخل کرنل و بازآرایی محاسبات (ترانهاده کردن $\mathcal{S}$ و $\mathcal{U}$ در ابتدا)، آن‌ها توانستند بهره‌وری بالای هسته‌های تنسور را حفظ کرده و همزمان دسترسی‌های غیرپیوسته (uncoalesced) به حافظه را حذف کنند.

بنچمارک و نتایج عملیاتی

محققان این کرنل‌ها را روی مجموعه متنوعی از مدل‌ها شامل Llama-7B، Llama-3.2-1B (که ۵۰٪ فشرده شده و به ۰.۶ میلیارد پارامتر رسیده بود)، GPT2-S، ViT-B و DiT-XL/2 آزمایش کردند. ارزیابی‌ها با دقت BF16 و استفاده از torch.compile() برای مقایسه‌های عادلانه با خط مبنا (baseline) انجام شد.

• Llama-7B (روی A40): پیاده‌سازی بهینه BLAST (روش ❺) شتاب ۳.۰۵ برابری در کل سیستم (end-to-end) نسبت به پیاده‌سازی اولیه BLAST ایجاد کرد.
• Llama-3.2-1B (روی Jetson): کرنل‌های بهینه‌شده شتاب ۳.۶۸ برابری را در کل سیستم فراهم کردند.
• DiT-XL/2 (روی Jetson): برای لایه $\text{QKV}_{proj}$، کرنل بهینه BLAST به اوج شتاب ۷.۱۵ برابر نسبت به خط مبنا دست یافت.
• عملکرد کلی: بهینه‌سازی ❺ در بیش از ۹۰٪ مواردی که خط مبنای BLAST شکست می‌خورد، توانست ۱.۱۳ تا ۳.۷۶ برابر سریع‌تر از مدل‌های متراکم (dense) عمل کند.

توازن میان دقت و سرعت

نکته حیاتی این است که این افزایش سرعت به بهای کیفیت مدل به‌دست نیامده است. جدول ۱ پژوهش نشان می‌دهد که BLAST به‌طور مداوم بالاترین دقت را در میان فرمت‌های فشرده در فاکتورهای فشرده‌سازی ۱.۸ تا ۳ برابر (CF 1.8x to 3x) به دست می‌آورد. برای مثال، در Llama-7B، روش BLAST در مقایسه با روش‌های کم‌رتبه سنتی، پرپلکسیتی (perplexity) برتر و دقت zero-shot بالاتری را حفظ می‌کند.

هنگام مقایسه پرپلکسیتی در برابر شتاب برای مدل‌های زبانی، کرنل‌های بهینه Monarch و BLAST یک جبهه پارتو (Pareto frontier) برتر نسبت به SVD سنتی ایجاد می‌کنند. این امر به توسعه‌دهندگان اجازه می‌دهد تا قدرت بیان مدل را حفظ کرده و در عین حال، مزایای سرعت وعده داده شده توسط کاهش تئوریک FLOPs را در عمل تجربه کنند.

تحلیل: چرخش به سوی کرنل‌های سفارشی

این پژوهش سیگنالی از یک تغییر حیاتی در نگاه ما به فشرده‌سازی مدل است. صنعت تا حد زیادی به بهینه‌سازی‌های سطح بالا در کامپایلر مانند torch.compile() تکیه کرده است، اما این تحقیق ثابت می‌کند که برای پراکندگی ساختاریافته (structured sparsity) و روش‌های کم‌رتبه بلوکی، ادغام در سطح کامپایلر کافی نیست. شکاف میان «وعده الگوریتمی» و «واقعیت سطح سیستم» برای کامپایلرهای عمومی بیش از حد زیاد است، به‌ویژه زمانی که با تانسورهای میانی $b \times n \times r$ در Monarch یا تانسورهای دوگانه در BLAST سروکار داریم.

با انتقال منطق محاسبات به درون کرنل‌های Triton، توسعه‌دهندگان می‌توانند کنترل دقیقی بر چیدمان حافظه و زمان‌بندی هسته‌های تنسور داشته باشند. این نشان می‌دهد که موج بعدی بهره‌وری هوش مصنوعی لبه (Edge AI) تنها از الگوریتم‌های هرس (pruning) بهتر نخواهد بود، بلکه حاصل یک «طراحی مشترک» (co-design) دقیق بین ساختار ریاضی وزن‌ها و اندازه دقیق کش L2 سخت‌افزار هدف است.

گام‌های بعدی

با کاربردی‌تر شدن روش‌های BLR، انتظار می‌رود ادغام کوانتش فعال‌سازها (activation quantization) برای کاهش بیشتر اندازه تانسورهای میانی صورت گیرد. مسیرهای آینده شامل بررسی کوانتش فعال‌سازهای میانی و گسترش استراتژی‌های ادغام برای پشتیبانی از پس‌انتشار (backpropagation) در طول آموزش است. توسعه‌دهندگان می‌توانند از کرنل‌های منتشرشده توسط این تیم برای پیاده‌سازی ساختارهای مشابه کم‌رتبه بلوکی در استقرار‌های محدود از نظر منابع خود استفاده کنند.