TileLang پیچیدگی هسته‌های GPU را تا ۹۰٪ کاهش داد

یک هسته سفارشی GPU که تنها ۵ میکروثانیه در زمان پاسخ‌دهی صرفه‌جویی کند، می‌تواند فشار پنج پردازنده H100 را در یک سیستم مقیاس‌پذیر به‌طور کامل بردارد. این واقعیت اقتصادی دلیل اصلی ظهور TileLang است؛ پروژه‌ای متن‌باز از پژوهشگران دانشگاه پکن و مایکروسافت که شیوه برخورد مهندسان با کارایی مدل‌های زبانی بزرگ (LLM) — مثل کتابخانه‌داری که میلیاردها صفحه را خوانده و حالا با همان لحن کتاب‌ها جواب می‌دهد — را تغییر می‌دهد.

کارایی هوش مصنوعی اکنون از معماری مدل‌ها فراتر رفته و به حوزه مهندسی هسته (Kernel Engineering) وارد شده است. در حالی که چارچوب‌هایی مثل PyTorch ضرب‌های ماتریسی استاندارد را مدیریت می‌کنند، در مواجهه با عملیات ترکیبی و غیرمعمول مدل‌های مدرن — مانند کوانتش وزن‌ها و مقیاس‌بندی آن‌ها در یک مرحله — دچار مشکل می‌شوند. همان‌طور که در تحلیل قبلی ما درباره ابزارهای بصری‌سازی کیفیت کد اشاره کردیم، TileLang اکنون شکاف اجرایی بین گراف‌های ریاضی و سخت‌افزار GPU را پر می‌کند.

به نقل از گزارشی در ۲۲ سپتامبر ۲۰۲۶، TileLang بر روی زیرساخت کامپایلر TVM بنا شده است. این زبان به توسعه‌دهندگان اجازه می‌دهد هسته‌های با کارایی بالا را با مدل برنامه‌نویسی کاشی‌بندی‌شده در پایتون بنویسند. با این روش، دیگر نیازی به درگیری‌های دستی با زبان CUDA C++ نیست، اما کنترل دقیق روی جای‌گذاری حافظه و موازی‌سازی حفظ می‌شود. این پروژه در ژانویه ۲۰۲۵ متن‌باز شد و نتایج آن در مقاله ICLR ۲۰۲۶ با عنوان «TileLang: پل میان برنامه‌پذیری و کارایی در هسته‌های عصبی مدرن» منتشر شد.

زمینه مهندسی هسته

برای درک TileLang باید فاصله بین گراف محاسباتی ریاضی و سخت‌افزار فیزیکی GPU را شناخت. یک مدل ممکن است عملیاتی ساده را تعریف کند، اما GPU به مسیر دقیق داده‌ها اهمیت می‌دهد: از حافظه HBM به L2، سپس به حافظه مشترک، ثبات‌ها و در نهایت هسته‌های تنسور / واحدهای ALU برداری.

مهندسی هسته دقیقاً در همین فاصله رخ می‌دهد. برای عملگرهای متداول، توسعه‌دهندگان از کتابخانه‌هایی مثل cuBLAS استفاده می‌کنند. اما استنتاج (Inference) — لحظه‌ای که مدل واقعاً جواب تولید می‌کند، شبیه خودِ آشپزی و نه دوره‌ی آموزش آشپز — در مدل‌های زبانی اغلب به توالی‌های غیرمعمولی نیاز دارد؛ مثلاً: کوانتش وزن‌ها $\rightarrow$ ضرب ماتریسی $\rightarrow$ مقیاس‌بندی $\rightarrow$ فعال‌سازی $\rightarrow$ یک تبدیل دیگر. یا در مکانیزم توجه: $Q \times QK^T \rightarrow$ softmax $\rightarrow$ $\times V$. معماری‌های جدیدتر، مانند Multi-Head Latent Attention در مدل DeepSeek، ترکیبات پیچیده‌تری از تصویرسازی‌ها (Projections)، الگوهای دسترسی به KV-cache، کاهش‌ها (Reductions) و جابجایی داده‌ها را معرفی می‌کنند.

این چالش منجر به تکامل ابزارهایی شد که ساختار را بدون نیاز به دستورات تک‌تک ماشین تعریف کنند. این مسیر از CUDA شروع شد، به کتابخانه‌های GPU رسید، از آنجا به TVM و کامپایلرهای تنسور رفت، سپس Triton و در نهایت به TileLang ختم شد.

نقاط عطف این تکامل عبارت‌اند از:

• ۲۰۱۷: معرفی معماری Volta انویدیا و هسته‌های تنسور که عملیات ماتریسی تخصصی را به یک ویژگی سخت‌افزاری درجه اول تبدیل کرد.
• ۲۰۱۸: معرفی TVM توسط Tianqi Chen و همکارانش به عنوان پشته‌ای برای نگاشت محاسبات تنسور روی سخت‌افزارهای متنوع.
• ۲۰۱۹: انتشار Triton توسط Philippe Tillet، H. T. Kung و David Cox که از «کاشی‌ها» به عنوان انتزاع مرکزی برای محاسبات شبکه عصبی استفاده کرد.
• ۲۰۲۲: نمایش FlashAttention توسط Tri Dao و همکارانش که ثابت کرد سازماندهی الگوریتمیک جابجایی حافظه می‌تواند کارایی را به‌طور رادیکال تغییر دهد.

سازوکار کاشی‌بندی

TileLang مفهوم میلیون‌ها رشته (Thread) مجزا در GPU را با تعداد کمی «کاشی» (Tile) جایگزین می‌کند که در سلسله‌مراتب حافظه حرکت می‌کنند. در یک ضرب ماتریسی ساده، هر عنصر خروجی نیاز به خواندن‌های مکرر از حافظه سراسری دارد که ترافیک عظیمی ایجاد می‌کند.

TileLang این مشکل را با بارگذاری بلوک‌های داده — یا همان کاشی‌ها — در حافظه مشترک حل می‌کند. برای مثال، یک بلوک رشته‌ای ممکن است یک کاشی ۱۲۸x۳۲ از ماتریس A و یک کاشی ۳۲x۱۲۸ از ماتریس B را بارگذاری کند تا یک کاشی ۱۲۸x۱۲۸ از ماتریس C تولید شود. این مقادیر بارها روی تراشه بازاستفاده می‌شوند و نیاز به مراجعه به حافظه پهنای‌باند بالا (HBM) را به‌شدت کاهش می‌دهند.

برای درک مقیاس این موضوع، دو ماتریس ۴۰۹۶x۴۰۹۶ با دقت FP16 را در نظر بگیرید. حجم محاسباتی ریاضی تقریباً ۱۳۷ میلیارد FLOP است ($2 \times 4096^3$). سه ماتریس در مجموع حدود ۹۶ مگابایت فضا اشغال می‌کنند (هر کدام ۳۲ مگابایت).

• پیاده‌سازی ساده: هر یک از ۱۶.۸ میلیون عنصر خروجی به ۴۰۹۶ مقدار از A و ۴۰۹۶ مقدار از B نیاز دارد که حدود ۲۷۵ گیگابایت ترافیک ورودی ایجاد می‌کند.
• پیاده‌سازی کاشی‌بندی‌شده: داده‌ها از حافظه روی تراشه بازاستفاده می‌شوند و ترافیک به عدد ایده‌آل ۹۶ مگابایت نزدیک می‌شود.

در یک پردازنده H100 SXM، که پهنای‌باند HBM آن ۳.۳۵ ترابایت بر ثانیه و توان محاسباتی هسته‌های تنسور آن ۱.۹۸ پتافلاپس (FP16) است، گلوگاه اصلی همیشه رساندن داده‌ها به واحدهای محاسباتی است، نه خودِ محاسبات. به همین دلیل بهینه‌سازی GPU یعنی جابجایی هر مقدار در کمترین تعداد دفعات ممکن. در اینجا کاشی هم واحد کار است و هم واحد جابجایی داده.

مدل برنامه‌نویسی و نحو

کدهای TileLang شبیه پایتون هستند اما جابجایی‌های سطح سخت‌افزار را توصیف می‌کنند. یک هسته ضرب ماتریسی از دکوراتور @tilelang.jit استفاده کرده و توابع را در بستر T.Kernel تعریف می‌کند. توسعه‌دهندگان از دستورات خاصی برای مدیریت حافظه استفاده می‌کنند:

• T.alloc_shared: تخصیص صریح فضا در حافظه مشترک (مثلاً A_shared = T.alloc_shared((block_M, block_K), dtype)).
• T.alloc_fragment: تخصیص فضا در ثبات‌ها برای جمع‌بندی مقادیر (مثلاً C_local = T.alloc_fragment((block_M, block_N), accum_dtype)).
• T.copy: جابجایی داده بین سطوح حافظه (مثلاً از حافظه سراسری به مشترک).
• T.gemm: اجرای ضرب ماتریسی روی کاشی‌ها با استفاده از دستورات هسته تنسور.
• T.Pipelined: هم‌پوشانی جابجایی داده با محاسبات برای حذف تأخیر، که اغلب از چندین مرحله استفاده می‌کند (مثلاً num_stages=3).

جداسازی «چه چیزی محاسبه شود» از «چگونه زمان‌بندی شود»، نوآوری اصلی این پروژه است. TileLang چهار بعد زمان‌بندی را مدیریت می‌کند:

• اتصال رشته‌ها (Thread Binding): نحوه نگاشت کار به رشته‌های GPU.
• چیدمان حافظه (Memory Layout): نحوه توزیع فیزیکی داده‌ها (که تحت تأثیر T.annotate_layout است).
• تنسورسازی (Tensorization): نگاشت عملیات به دستورات هسته تنسور.
• خط لوله (Pipeline): کنترل هم‌پوشانی جابجایی داده و محاسبات از طریق T.Pipelined.

توسعه‌دهنده جریان داده ریاضی را تعریف می‌کند و کامپایلر جزئیات مکانیکی نگاشت رشته‌ها و استنتاج چیدمان را بر عهده می‌گیرد. این موضوع حیاتی است زیرا زمان‌بندی که روی A100 جواب می‌دهد، ممکن است روی H100 متفاوت عمل کند یا برای سخت‌افزارهای AMD به پیاده‌سازی متفاوتی نیاز داشته باشد.

کاربردهای واقعی در مدل‌های زبانی

TileLang برای هسته‌های پیچیده‌ای طراحی شده که کارایی LLMهای مدرن را تعیین می‌کنند. طبق مقاله ICLR ۲۰۲۶، کاربردهای کلیدی آن عبارت‌اند از:

پیاده‌سازی توجه برق‌آسا (FlashAttention): توجه معمولی باعث ایجاد یک ماتریس $n \times n$ می‌شود که باعث رشد مربعی حافظه می‌گردد. FlashAttention با استفاده از فرمول‌بندی آگاه از ورودی/خروجی (IO-aware)، بلوک‌های Q و K را پردازش کرده و مقادیر میانی را روی تراشه نگه می‌دارد. TileLang می‌تواند زمان‌بندی‌های خط لوله‌ای به پیچیدگی FlashAttention-3 را پیاده کند. در ارزیابی‌ها، پیاده‌سازی TileLang در بارهای کاری مورد آزمایش، از Triton و PyTorch پیشی گرفته و در توالی‌های طولانی‌تر، به عملکرد FlashAttention-3 نزدیک ماند.

استنتاج کوانتیده: برای مدل‌هایی با وزن‌های INT4 یا FP4، TileLang امکان «GEMM فقط-وزن» را فراهم می‌کند. به‌جای اجرای دو هسته جدا برای کوانتش و ضرب — که ترافیک حافظه اضافی ایجاد می‌کند — TileLang وزن‌ها را مستقیماً در یک کاشی محلی کوانتش کرده و به هسته‌های تنسور می‌فرستد. این جریان به صورت: داده‌های فشرده $\rightarrow$ کاشی $\rightarrow$ کوانتش $\rightarrow$ کاشی $\rightarrow$ GEMM هسته تنسور است.

توجه به سبک DeepSeek: این پروژه شامل پیاده‌سازی FlashMLA (Multi-Head Latent Attention) است. این هسته مدیریت کاشی‌های Q، K و مؤلفه‌های موقعیتی را به همراه جمع‌بندی امتیازات، کاهش‌ها (T.reduce_max, T.reduce_sum) و خط لوله‌سازی بر عهده دارد. این هسته پیچیده در کمتر از ۸۰ خط پایتون نوشته شده است، که نشان می‌دهد هسته همزمان یک محاسبه، یک استراتژی مدیریت حافظه و یک نگاشت سخت‌افزاری است.

نقش کامپایلر

TileLang صرفاً یک پوشش برای CUDA نیست، بلکه یک خط لوله کامپایلر است. این ابزار کد پایتون را از طریق یک AST به نمایش میانی TVM تبدیل کرده و پس از اعمال گذرهای بهینه‌سازی، کد مخصوص سخت‌افزار هدف را تولید می‌کند. این سیستم از طیف وسیعی از بک‌اِندها شامل CUDA، HIP، Metal، WebGPU و اجرای CPU پشتیبانی می‌کند.

یکی از ویژگی‌های کلیدی، استنتاج چیدمان (Layout Inference) است. وقتی توسعه‌دهنده یک کاشی منطقی ۱۲۸x۱۲۸ تعریف می‌کند، کامپایلر تعیین می‌کند که این کاشی چگونه به‌طور فیزیکی بین وارپ‌ها (Warps) و قطعات ثبات (Register Fragments) توزیع شود. این مسیر به صورت: کاشی منطقی $\rightarrow$ تقسیم $\rightarrow$ وارپ‌ها/رشته‌ها $\rightarrow$ قطعات ثبات $\rightarrow$ چیدمان دستور هسته تنسور است.

پژوهش ICLR ۲۰۲۶ دو قابلیت پیشرفته را معرفی می‌کند:

• استنتاج کاشی (Tile Inference): استفاده از ساختار یک برنامه کاشی ترکیبی برای استنتاج اطلاعات پیکربندی مفقود.
• پیشنهاد کاشی (Tile Recommendation): استفاده از اطلاعات سخت‌افزاری و اکتشافی (Heuristics) برای پیشنهاد بهینه‌ترین پیکربندی‌ها.

در آزمایش‌ها، این رویکرد حجم کد را در مقایسه با پیاده‌سازی دستی CUDA تا ۹۰٪ کاهش داد. هدف جایگزینی مهندس کارایی نیست، بلکه انتقال تمرکز او از سطح دستورات به سطح الگوریتم است.

اقتصاد بهینه‌سازی

طبق قانون آمدال، ارزش یک هسته TileLang به سهم آن از زمان کل اجرا بستگی دارد. اگر یک هسته خاص ۴۰٪ از زمان یک سرویس استنتاج را مصرف کند، دو برابر کردن سرعت آن هسته، زمان کل برنامه را ۱.۲۵ برابر بهبود می‌بخشد (کاهش ۱۰۰ واحد زمانی به ۸۰ واحد).

برای سیستم‌های تولیدی، فرمول ساده است: $\text{ارزش} = \text{تعداد فراخوانی‌ها} \times \text{زمان ذخیره شده} \times \text{هزینه محاسبات}$. وقتی یک هسته میلیون‌ها بار در ثانیه فراخوانی می‌شود، ذخیره چند میکروثانیه مستقیماً به کاهش تعداد GPUهای مورد نیاز در یک خوشه منجر می‌شود. برای مثال، ذخیره ۵ میکروثانیه در ۱ میلیون فراخوانی در ثانیه، معادل ۵ ثانیه زمان محاسباتی GPU در هر ثانیه است، که یعنی آزاد شدن پنج پردازنده H100 که به‌طور مداوم اشغال شده بودند.

گردش کار توسعه‌دهنده

برای کسانی که می‌خواهند TileLang را به کار بگیرند، مسیر پیشنهادی تکرارپذیر است. توسعه‌دهندگان باید با یک عملگر PyTorch شروع کنند، آن را پروفایل کنند تا گلوگاه را بیابند و تنها در آن صورت به سراغ هسته سفارشی Triton یا TileLang بروند.

TileLang زمانی جذاب می‌شود که گلوگاه شامل موارد زیر باشد:

• ترکیب غیرمعمول عملیات (Fusion)
• محاسبات کوانتیده (INT4/FP4)
• مکانیزم‌های توجه سفارشی
• کاهش‌های (Reductions) غیرمعمول
• جابجایی‌های گران‌قیمت حافظه
• نیاز به جای‌گذاری صریح در حافظه مشترک یا ثبات‌ها
• الزامات خط لوله‌ای مخصوص معماری
• تمایل به هدف قرار دادن چندین بک‌اِند شتاب‌دهنده (CUDA, HIP, Metal)

این تغییر نشان‌دهنده یک روند کلی در هوش مصنوعی است: مهندسی کارایی در حال تبدیل شدن به رشته «طراحی جابجایی داده» است. واحد بنیادی بهینه‌سازی دیگر عدد (Scalar) نیست، بلکه کاشی است. مسئله اصلی دیگر فقط معادله $C = A @ B$ نیست، بلکه تصمیم‌گیری درباره این است که کدام تکه از A و B در HBM بماند، چگونه به حافظه مشترک منتقل شود و چگونه توسط رشته‌ها بازاستفاده شود. پرسش‌های محوری این‌ها هستند: کدام A؟ کدام B؟ کجا زندگی می‌کنند؟ چه زمانی بارگذاری می‌شوند؟ چه کسی آن‌ها را بارگذاری می‌کند؟ چه کسی بازاستفاده می‌کند؟ کدام رشته‌ها مالک آن‌ها هستند؟ کدام دستور آن‌ها را محاسبه می‌کند؟ آیا کاشی بعدی می‌تواند در حین محاسبه این یکی بارگذاری شود؟

نتیجه‌گیری

تکامل از CUDA به TVM، Triton و TileLang بازتاب‌دهنده تنش همیشگی در برنامه‌نویسی سیستم‌ها بین بهره‌وری سطح بالا و کارایی سطح پایین است. TileLang مرز انتزاع را به سمت بالا می‌برد بدون اینکه تظاهر کند جزئیات سخت‌افزاری ناپدید شده‌اند. برای توسعه‌دهندگان LLM، این موضوع ضروری است زیرا هسته‌های امروزی به‌طور فزاینده‌ای ترکیبی و وابسته به معماری هستند. وقتی یک هسته توجه را به عنوان جریانی از کاشی‌ها می‌بینید که بین HBM، حافظه مشترک، ثبات‌ها و هسته‌های تنسور حرکت می‌کنند، استدلال درباره کد بسیار آسان‌تر می‌شود. وقتی دفعه بعد یک LLM را پروفایل کردید و متوجه شدید یک هسته سفارشی کوچک مقدار غافلگیرکننده‌ای از زمان اجرا را می‌گیرد، انتخاب بین CUDA، Triton یا TileLang به این بستگی دارد که برای دستیابی به کارایی لازم، به چه میزان کنترل صریح روی جریان داده نیاز دارید.

گام بعدی شما

• اگر از مدل‌های کوانتیده استفاده می‌کنید، بررسی کنید که آیا عملیات Dequantization و GEMM شما در دو هسته جداگانه اجرا می‌شوند یا به صورت Fused.
• برای کاهش هزینه‌های استنتاج در مقیاس بالا، پروفایلینگ دقیق هسته‌های سفارشی را جایگزین بهینه‌سازی‌های کلی معماری کنید.
• مستندات TileLang را برای پیاده‌سازی مکانیزم‌های Attention سفارشی مطالعه کنید تا حجم کد خود را کاهش دهید.

اما داستان سخت‌افزاری این تحول حتی شگفت‌انگیزتر است — به تحلیل ما درباره تراشه‌های Blackwell مراجعه کنید.