تصور کنید به جای مدیریت تک‌تک رشته‌ها، کل بلوک‌های داده را سازمان‌دهی کنید؛ این دقیقاً همان نیازی است که برای نوشتن هسته‌های GPU با کارایی بالا دارید. از طریق cuTile و Triton، مدل برنامه‌نویسی کاشی‌محور انویدیا به توسعه‌دهندگان اجازه می‌دهد یک کاشی کامل از داده را بارگذاری کرده، روی آن به‌طور بهینه محاسبه کنند و نتیجه را در یک گذر واحد ذخیره کنند. این رویکرد به‌طور مؤثر ناکارآمدی‌های ذاتی موجود در الگوهای کلاسیک SIMT (یک دستورالعمل، چندین رشته) را دور می‌زند.
بارهای کاری مدرن هوش مصنوعی به بهره‌وری شدید حافظه نیاز دارند تا از «دیوار حافظه» (memory wall) عبور کنند؛ وضعیتی که در آن GPU زمان بیشتری را صرف جابه‌جایی داده‌ها می‌کند تا انجام محاسبات. همان‌طور که در تحلیل‌های پیشین ما درباره‌ی بهینه‌سازی لایه‌های حافظه در مدل‌های زبانی اشاره کردیم، این چرخش به سمت کاشی‌بندی برای مقیاس‌پذیری مدل‌ها حیاتی است، زیرا عملیاتی را می‌طلبد که بتوانند چندین مرحله ریاضی را در یک سفر حافظه ادغام کنند. این نیاز به بهینه‌سازی دقیق سخت‌افزاری، در واقع پاسخی به چالش‌های زیرساختی و پراکندگی منابع GPU است که می‌تواند مانع از رسیدن به پتانسیل کامل معماری مدل‌ها شود. برای توسعه‌دهندگانی که از Google Colab یا ایستگاه‌های کاری محلی استفاده می‌کنند، این انتقال معمولاً در سه سطح رخ می‌دهد: GPUهای سطح بالای Ampere به بالا با CUDA 13.1، GPUهای استاندارد که از Triton استفاده می‌کنند، یا محاسبات مرجع مبتنی بر CPU برای تست صحت (correctness testing).
مکانیسم مدل کاشی
در یک ساختار سنتی CUDA، شما کدی برای یک رشته (thread) می‌نویسید. آن رشته یک شاخص سراسری را محاسبه می‌کند، محدوده‌ها را بررسی می‌کند و تنها به یک عنصر دسترسی می‌یابد. منطق این کار اساساً به شکل C[i] = A[i] + B[i] است. در مقابل، مدل کاشی بر روی بلوکی تمرکز می‌کند که مالک یک زیرماتریس یا بردار خاص است (به عنوان مثال، ۱۰۲۴ عنصر یا یک زیرماتریس ۱۲۸x۱۲۸).
- پایه-توابع cuTile مانند
ct.bid(0)برای شناسه‌های بلوک،ct.loadوct.storeجابه‌جایی این بلوک‌ها را مدیریت می‌کنند. این مدل همچنین از عملیات مستقیم ماتریسی مانندa @ bو اجرای هسته (kernel launching) از طریقct.launchپشتیبانی می‌کند. - Triton برای رسیدن به نتیجه منطقی مشابه، از رویکردی مبتنی بر شبکه (grid) با استفاده از
tl.program_id،tl.load،tl.storeوtl.dotبهره می‌برد. - کامپایلر مدیریت نگاشت این کاشی‌ها روی هسته‌های تانسوری (Tensor Cores) فیزیکی را بر عهده می‌گیرد و بار دستی برنامه‌نویس را برای مدیریت سخت‌افزار حذف می‌کند.
بررسی محیط و انتخاب بک‌اند
برای اطمینان از اجرای کد روی سخت‌افزارهای متنوع، گردش کار TileGym با یک بررسی جامع محیطی آغاز می‌شود. سیستم ابتدا با استفاده از تابع torch.cuda.is_available() وجود CUDA را چک می‌کند. در صورت یافتن GPU، سیستم نام دستگاه و قابلیت محاسباتی (Compute Capability یا به اختصار CC) را استخراج می‌کند تا سطح پشتیبانی سخت‌افزاری را تعیین کند.
منطق سازگاری بک‌اند
انتخاب بک‌اند اجرایی از یک سلسله‌مراتب سخت‌گیرانه بر اساس نسخه‌های سخت‌افزاری و نرم‌افزاری پیروی می‌کند:
- بک‌اند cuTile: این گزینه «استاندارد طلایی» است. برای استفاده از آن، دستگاه باید قابلیت محاسباتی ۸.۰ یا بالاتر داشته باشد (معماری‌های Ampere، Ada یا Blackwell) و نسخه CUDA نصب شده باید ۱۳.۱ یا جدیدتر باشد. اگر این شرایط برقرار باشد، سیستم تلاش می‌کند
cuda.tileرا وارد کند و در صورت نیاز، بسته هایcuda-tile[tileiras]وcupy-cuda13xرا نصب نماید. - بک‌اند Triton: این مسیر استاندارد در محیط Colab است. اگر GPU شناسایی شده از نوع T4 باشد (که CC آن کمتر از ۸.۰ است) یا اگر نسخه CUDA پایین‌تر از ۱۳.۱ باشد، سیستم به Triton باز می‌گردد. این بک‌اند مدل برنامه‌نویسی کاشی‌محور مشابهی را ارائه می‌دهد اما با طیف وسیع‌تری از سخت‌افزارها سازگار است.
- بک‌اند Torch/CPU: در صورتی که هیچ GPU با قابلیت CUDA یافت نشود، سیستم به طور پیش‌فرض به عملیات استاندارد PyTorch روی CPU می‌رود تا محاسبات ریاضی مرجع و تست‌های صحت امکان‌پذیر باشد و جریان توسعه متوقف نشود.
جزئیات نیازمندی‌های بک‌اند
در طول فرآیند بررسی محیط، دلایل خاص شکست‌ها در لاگ‌ها ثبت می‌شود تا کاربر را راهنمایی کند:
- اگر مقدار
HAS_CUDAغلط (false) باشد، سیستم عبارت "no CUDA GPU" را ثبت می‌کند. - اگر CC کمتر از ۸.۰ باشد، سیستم ذکر می‌کند که سخت‌افزارهای Turing یا T4 برای اجرای واقعی cuTile پشتیبانی نمی‌شوند.
- اگر نسخه CUDA پایین‌تر از ۱۳.۱ باشد، سیستم هشدار می‌دهد که پیش‌نیاز
tileirasبرآورده نشده است.
پیاده‌سازی هسته‌های اصلی
به نقل از راهنمای منتشر شده توسط Marktechpost، کاربرد عملی این مدل با جمع برداری ساده شروع می‌شود. به جای استفاده از یک حلقه عنصر-به-عنصر، هسته یک کاشی با اندازه ۱۰۲۴ را بارگذاری کرده، آن‌ها را جمع کرده و سپس نتایج را بازمی‌گرداند. پیاده‌سازی Triton در اینجا از tl.program_id(0) و tl.arange(0, BLOCK) برای ایجاد آفست‌ها و یک ماسک (mask) برای مدیریت شرایط مرزی استفاده می‌کند؛ این کار برای مواردی است که تعداد کل عناصر n مضرب اندازه بلوک نباشد.
یک مثال کلیدی و پیشرفته، هسته Fused GELU است. در PyTorch استاندارد، عملیاتی مانند GELU(x * w + b) به سه گذر حافظه جداگانه نیاز دارد (یک بار برای ضرب و جمع، و یک بار برای تابع فعال‌ساز). یک هسته کاشی‌محور این مراحل را در یک گذر واحد ادغام (fuse) می‌کند و ترافیک بین حافظه سراسری GPU و ثبات‌ها (registers) را به‌شدت کاهش می‌دهد. این هسته فرمول تقریب tanh-GELU را پیاده می‌کند:
$0.5 * h * (1.0 + tanh(0.7978845608028654 * (h + 0.044715 * h^3)))$
که در آن $h = x * w + b$ است.
جزئیات پیاده‌سازی GELU
- ورودی‌ها: هسته چهار اشاره‌گر (pointer) دریافت می‌کند:
x_ptr،w_ptr،b_ptrوo_ptr. - ثابت‌ها: برای این تقریب ریاضی، از ثابت دقیق
c = 0.7978845608028654استفاده می‌شود. - ادغام (Fusion): با انجام عملیات ضرب-جمع و فعال‌ساز در یک هسته واحد، سیستم از نوشتن نتایج میانی در VRAM جلوگیری می‌کند که منجر به افزایش چشمگیر سرعت می‌شود.
برای پیاده‌سازی سافت‌مکس (Softmax) سطری، مدل کاشی از کاهش‌های عددی پایدار (numerically stable reductions) استفاده می‌کند. هسته در هر برنامه (program ID) یک سطر را پردازش می‌کند. ابتدا مقدار حداکثری را در کل کاشی (با استفاده از tl.max با axis=0) می‌یابد تا از سرریز (overflow) هنگام محاسبه توابع نمایی جلوگیری کند. سپس مجموع توابع نمایی (مخرج کسر) را محاسبه کرده و تقسیم نهایی را انجام می‌دهد، در حالی که داده‌ها در تمام این مدت در سطح بلوک محلی باقی می‌مانند.
مکانیسم Softmax
- پایداری عددی: هسته مقدار حداکثر را از ورودی کم می‌کند:
x = x - tl.max(x, axis=0). - کاشی‌بندی: اندازه بلوک توسط
triton.next_power_of_2(ncols)تعیین می‌شود تا پوشش کامل سطر تضمین گردد. - ماسک‌گذاری: برای مدیریت ابعادی که توان دو نیستند، یک ماسک
cols < n_colsاعمال می‌شود و برای عملیات max، مقدارother=-float("inf")استفاده می‌شود تا عناصر خارج از محدوده تأثیری در نتیجه نداشته باشند.
مقیاس‌پذیری به Matmul و Flash Attention
ضرب ماتریسی ستون فقرات مدل‌های زبانی بزرگ (LLMs) است. رویکرد کاشی‌بندی این مشکل را با استفاده از یک «حلقه K» برای تجمع (accumulation) حل می‌کند. به جای محاسبه کل ماتریس خروجی به صورت یکجا، GPU بلوک‌های کوچکی را بارگذاری می‌کند (معمولاً BM=64، BN=64 و BK=32). سپس ضرب داخلی را از طریق هسته‌های تانسوری با استفاده از tl.dot انجام داده، نتایج را در یک ثبات محلی float32 جمع‌آوری می‌کند و در نهایت پس از اتمام حلقه K، نتیجه را در ماتریس خروجی C ذخیره می‌کند.
پیشرفته‌ترین کاربرد این مدل، هسته توجه برق‌آسا (Flash Attention) است. این پیاده‌سازی از تکنیک «سافت‌مکس آنلاین» استفاده می‌کند تا از ایجاد و ذخیره ماتریس کامل $L \times L$ در حافظه جلوگیری کند. این هسته روی تانسورهای Query (Q)، Key (K) و Value (V) با اندازه بلوک (BL) برابر ۶۴ و بُعد (BD) برابر با بُعد سر (head dimension یا D) عمل می‌کند.
مکانیسم Flash Attention
هسته Flash Attention کارایی خود را از طریق گام‌های زیر به دست می‌آورد:
- سافت‌مکس افزایشی: یک مقدار حداکثر جاری (
m_i) و یک مجموع نمایی جاری (l_i) برای هر سطر نگه می‌دارد. این مقادیر به ترتیب با-float("inf")و0.0مقداردهی اولیه می‌شوند. - تکرار کاشی‌بندی: طول توالی $L$ را در بلوک‌های $BL$ پیمایش کرده و برش‌های تانسورهای K و V را با استفاده از
tl.loadو ماسک‌های مرزی بارگذاری می‌کند. - مقیاس‌بندی مجدد (Rescaling): هرگاه حداکثر جدیدی یافت شود، تجمع قبلی (
acc) را با استفاده از فاکتور آلفا $\exp(m_i - m_{ij})$ بازتنظیم می‌کند تا صحت عددی محاسبات حفظ شود. - نرمال‌سازی نهایی: پیش از آنکه نتایج با استفاده از
tl.storeو بررسی ماسکoffs_l[:, None] < Lدر تانسور خروجیOذخیره شوند، مجموع تجمع‌شده بر مقدار نهاییl_iتقسیم می‌شود.
بنچ‌مارک و اعتبارسنجی
برای اطمینان از اینکه این هسته‌های سفارشی به درستی کار می‌کنند، باید با استفاده از یک ابزار check در برابر مراجع PyTorch اعتبارسنجی شوند. این ابزار حداکثر تفاوت مطلق (md) بین خروجی مدل کاشی-محور و خروجی PyTorch را محاسبه می‌کند.
- جمع برداری و GELU: این موارد با تلورانس استاندارد
atolوrtolبرابر با 1e-2 بررسی می‌شوند. - ضرب ماتریسی (Matmul): به دلیل ماهیت تجمع اعداد اعشاری و استفاده از هسته‌های تانسوری، Matmul با تلورانس آزادتر
atol=1e-1وrtol=1e-1برای ورودی‌هایی با اندازه ۱۰۲۴x۱۰۲۴ تأیید می‌شود. - Flash Attention: این هسته در برابر
torch.nn.functional.scaled_dot_product_attentionبا طول توالی ۵۱۲، بُعد سر ۶۴ و ۸ سر، با تلورانس 2e-2 اعتبارسنجی می‌گردد.
عملکرد با یک تابع bench اندازه‌گیری می‌شود که شامل یک فاز گرم‌کردن (۱۰ تکرار) و یک فاز زمان‌بندی شده (۵۰ تکرار) است تا اطمینان حاصل شود که GPU در حداکثر ساعت کاری خود است و با torch.cuda.synchronize() کاملاً همگام شده است. برای Matmul، توان عملیاتی بر حسب TFLOP/s با فرمول $2 \times M \times N \times K / (time \times 10^{-3}) / 10^{12}$ محاسبه می‌شود. برای تست جمع برداری، از برداری با اندازه $2^{20}$ عنصر استفاده می‌گردد.
نیازمندی‌های سخت‌افزاری cuTile
برای اجرای واقعی هسته‌های cuTile مطابق با کدهای منبع، کاربران به پارامترهای محیطی خاصی نیاز دارند:
- GPU: معماری Ampere، Ada یا Blackwell (قابلیت محاسباتی ۸.۰+).
- نرم‌افزار: CUDA Toolkit نسخه ۱۳.۱ یا بالاتر.
- کتابخانه‌ها: نصب
cuda-tile[tileiras]وcupy-cuda13xالزامی است. برای نصب کامل این مجموعه می‌توان از دستورpip install tilegym[tileiras]استفاده کرد.
برای کسانی که سخت‌افزار قدیمی‌تری مانند NVIDIA T4 دارند، بک‌اند Triton تجربه‌ای تقریباً یکسان را فراهم می‌کند. این موضوع تضمین می‌کند که منطق محاسبات کاشی-محور (بارگذاری بلوک، محاسبه و ذخیره) حتی بدون داشتن جدیدترین سخت‌افزارهای انویدیا در دسترس باشد.
این تکامل در برنامه‌نویسی GPU نشان‌دهنده یک روند گسترده‌تر است: با رشد پیچیدگی معماری‌های AI، لایه انتزاع باید از «اسکالارهای تک‌عنصری» به «بلوک‌های داده» حرکت کند تا بتواند با توان سخت‌افزار همگام شود. با تمرکز بر کاشی (Tile)، توسعه‌دهندگان می‌توانند هسته‌هایی بنویسند که به سخت‌افزار نزدیک‌تر هستند، بدون اینکه نیاز داشته باشند هر یک از Warpهای CUDA را به صورت دستی مدیریت کنند.
گام بعدی شما
- اگر از GPUهای سری RTX 30 یا 40 استفاده می‌کنید، کتابخانه
tilegymرا برای تجربه برنامه‌نویسی کاشی‌محور نصب کنید. - کدهای Triton خود را با استراتژی «ادغام هسته» (Kernel Fusion) بازبینی کنید تا دفعات دسترسی به VRAM را کاهش دهید.
- برای اعتبارسنجی هسته‌های سفارشی، از متد مقایسه تفاوت مطلق (Absolute Difference) با مراجع PyTorch استفاده کنید.
اما داستان سخت‌افزاری این تحول حتی شگفت‌انگیزتر است — به تحلیل ما درباره‌ی تراشه‌های Blackwell مراجعه کنید.
گفتگو