یک خط کد در PyTorch می‌تواند مدل شما را شتاب ببخشد یا عملکرد آن را تا ۴ برابر تخریب کند. این تضاد تکان‌دهنده در سری تحلیل‌های پروفایلینگ شرکت Hugging Face که در ۱۰ ژوئیه ۲۰۲۶ منتشر شد، آشکار شد؛ جایی که مشخص شد انتخاب بک‌اند برای scaled_dot_product_attention (SDPA) به‌طور کلی نحوه مدیریت مکانیزم توجه (Attention) توسط GPU را تغییر می‌دهد.
برای توسعه‌دهندگانی که مدل‌های زبانی بزرگ (LLM) — مثل کتابخانه‌داری که میلیاردها صفحه را خوانده و حالا با همان لحن جواب می‌دهد — یا مدل‌های انتشار را می‌سازند، توجه گلوگاه اصلی است چون پیچیدگی زمانی آن به‌صورت درجه‌دوم رشد می‌کند. اگرچه APIهای سطح بالا این پیچیدگی را می‌پوشانند، اما اجرای سخت‌افزاری در بک‌اندهای مختلف کاملاً متفاوت است. این موضوع صرفاً مربوط به کارایی نرم‌افزاری نیست، بلکه به نحوه جابه‌جایی داده‌ها بین حافظه پهنای‌باند بالا (HBM) و رجیسترهای داخلی تراشه بازمی‌گردد. طبق گزارش Hugging Face، این آزمایش‌ها روی یک GPU مدل NVIDIA A100-SXM4-80GB انجام شده است.
همان‌طور که در تحلیل‌های قبلی ما درباره بهینه‌سازی حافظه در مدل‌های ترنسفورمر اشاره کردیم، مدیریت ترافیک داده بین VRAM و هسته‌های پردازشی، تعیین‌کننده نهایی سرعت است.
پیاده‌سازی ساده و هزینه پنهان حافظه
ساخت مکانیزم توجه از قطعات اولیه (مانند ضرب ماتریسی، مقیاس‌بندی، ماسک‌گذاری و سافت‌مکس) یک نشت عملکردی جدی را فاش می‌کند. در یک پیاده‌سازی ساده، PyTorch ابتدا امتیازها را محاسبه کرده، آن‌ها را مقیاس می‌کند و سپس یک ماسک علی (Causal Mask) اعمال می‌کند.
بر اساس مستندات تحلیل شده، استفاده از دستور masked_fill باعث ایجاد یک هسته کپی حافظه (Memcpy) پنهان در GPU می‌شود. دلیل این اتفاق آن است که PyTorch اغلب یک کپی از تنسور می‌گیرد و عملیات را روی کپی انجام می‌دهد، نه روی نسخه اصلی.
با تغییر ساده به نسخه در-جا یا in-place (یعنی masked_fill_ که با یک خط تیره در پایان شناخته می‌شود)، هسته Memcpy کاملاً حذف می‌شود. ردپای پروفایلر نشان می‌دهد که نسخه in-place عملیات‌های CPU بسیار کمتری را در مرحله ماسک‌گذاری اجرا می‌کند.
این تغییر کوچک در هر لایه تکرار می‌شود و در یک مدل ترنسفورمر مدرن با ده‌ها لایه، مقدار قابل‌توجیهی از زمان و حافظه را ذخیره می‌کند. این روش به شرطی ایمن است که برنامه تحت torch.no_grad() اجرا شود؛ زیرا در این حالت نیازی به ذخیره مقادیر اصلی برای پاس بازگشتی (Backward Pass) نیست و بازنویسی داده‌ها مشکلی ایجاد نمی‌کند.
بک‌اند Math: تله‌ای برای عملکرد
وقتی توسعه‌دهندگان از F.scaled_dot_product_attention با تثبیت روی بک‌اند 'math' استفاده می‌کنند، با نتیجه‌ای شوکه کننده روبرو می‌شوند. به‌رغم اینکه فقط یک خط کد است، بک‌اند math حدود ۳.۷ برابر کندتر از پیاده‌سازی ساده in-place عمل می‌کند. در جدول پروفایلر، میانگین زمان CUDA برای عملیات *_fwd از ۱.۹۵۵ میلی‌ثانیه به ۷.۲۳۹ میلی‌ثانیه جهش می‌کند.
پروفایلر فاش می‌کند که بک‌اند math در هر پاس پیشرو ۲۰ هسته GPU را فراخوانی می‌کند، در حالی که نسخه ساده تنها ۵ تا ۶ هسته دارد. این سقوط عملکرد به سه دلیل رخ می‌دهد:
۱. خالی ماندن هسته‌های تنسور (Tensor Cores)
- ارتقای دقت به FP32: بک‌اند math برای دقت عددی بیشتر، تنسورها را به FP32 ارتقا می‌دهد. این کار حجم داده‌های جابه‌جایی را حتی برای ورودی‌های bf16 دو برابر می‌کند.
- بازگشت به هسته‌های CUDA: به دلیل تبدیل به FP32، مدل از هسته‌های تخصصی Tensor Core عبور کرده و به هسته‌های general-purpose قدیمی‌تر بازمی‌گردد. نام هسته‌ها در ردپا به‌جای امضاهای Tensor Core، عبارت
sgemmرا نشان می‌دهد.
۲. تجسم ماسک (Mask Materialization)
- ساخت لحظه‌ای: پرچم
is_causal=Trueکار ماسک‌گذاری را حذف نمی‌کند، بلکه فقط جابه‌جا می‌کند. بک‌اند math در هر فراخوانی، ماسک را از نو می‌سازد. - خط لوله CPU: در مسیر CPU، توالی دستورات
aten::onesوaten::trilبرای ساخت یک ماتریس مثلثی پایین دیده می‌شود که در نهایت به یک فشار پردازشی روی GPU منجر می‌گردد.
۳. سافت‌مکس ایمن (Safe Softmax)
- جلوگیری از NaN: برای جلوگیری از تولید مقادیر نامعتبر (NaN) در ردیف‌های کاملاً ماسک‌شده، بک‌اند math از
aten::_safe_softmaxاستفاده می‌کند که باعث اضافه شدن هسته‌های پردازشی اضافی در ردپا می‌شود.
هسته‌های ادغام‌شده: بک‌اندهای Efficient و Flash
برای فرار از پراکندگی ۲۰ هسته‌ای در بک‌اند math، PyTorch از بک‌اندهای بهینه مانند FLASH_ATTENTION و EFFICIENT_ATTENTION استفاده می‌کند. این‌ها تمام عملیات‌های اولیه را در یک تک‌هسته (Fused Kernel) ادغام می‌کنند.
بک‌اند 'efficient' (مبتنی بر xformers) از هسته fmha_cutlassF استفاده می‌کند که روی CUTLASS ساخته شده، در حالت bfloat16 اجرا می‌شود و داده‌ها را در رجیسترهای داخلی نگه می‌دارد تا سرعت در GPUهای نسل Ampere به حداکثر برسد.
توجه برق‌آسا (FlashAttention-2) یک گام فراتر می‌رود و ترافیک HBM را هدف قرار می‌دهد. هزینه اصلی توجه نه ضرب ماتریسی، بلکه ترافیک رفت‌وبرگشت مداوم به حافظه برای خواندن و نوشتن ماتریس امتیازات است. برای یک توالی ۴۰۹۶، این یعنی جابه‌جایی حدود ۱۶ میلیون عدد برای هر سر توجه.
Flash از ترفندی به نام «سافت‌مکس آنلاین» استفاده می‌کند تا روی K و V به‌صورت تکه‌های کوچک (Tiles) حرکت کند. در این حالت، ماتریس کامل امتیازات هرگز در HBM نوشته نمی‌شود و فقط روی تراشه می‌ماند.
ردپای «غلط» پروفایلر در Flash
به‌طور جالب، پروفایلر اشغال حافظه (Occupancy) پایینی (حدود ۱۳٪) را برای FlashAttention گزارش می‌کند. این به معنای بهینه‌سازی差 نیست. Flash به‌جای نگه داشتن تعداد زیادی Warp برای پنهان کردن تأخیر، بر بازاستفاده‌ی خام از داده‌ها تمرکز دارد.
- سنگینی منابع: Flash مقدار عظیمی از حافظه مشترک (Shared Memory) و رجیسترهای هر رشته را اشغال می‌کند.
- تعداد Warps: در یک SM نسل Ampere، به‌دلیل مصرف بالای رجیسترها، تنها ۲ بلوک جای می‌گیرند که منجر به عدد ۱۳٪ می‌شود.
رویکرد cuDNN: تولید لحظه‌ای هسته
برخلاف Flash، بک‌اند NVIDIA cuDNN هسته‌هایی را در زمان اجرا (Runtime) متناسب با ابعاد مسئله تولید می‌کند. نام‌های طولانی هسته‌ها در ردپا (مانند cudnn_generated_...) گواه این رویکرد است.
ویژگی‌های کلیدی cuDNN عبارتند از:
- تنظیم Knob: بر اساس ابعاد تنسور، یکی از پیکربندی‌های پیش‌تنظیم (Knobs) را انتخاب می‌کند.
- چیدمان مستقیم: داده‌ها را مستقیماً با ساختار [B, H, S, D] مصرف کرده و چهار عملیات
aten::transposeرا که در Flash دیده می‌شود، حذف می‌کند. - فراخوانی سطح درایور: از
cuLaunchKernelExاستفاده می‌کند که باعث می‌شود CUPTI نتواند اشغال حافظه را به‌درستی اندازه‌گیری کند و گاهی عدد ۰٪ را گزارش دهد.
این انعطاف‌پذیری هزینه‌ای روی CPU دارد. طبق داده‌های پروفایلر، cuDNN حدود ۲۱۴ میکروثانیه را صرف انتخاب بهترین Knob می‌کند، در حالی که این زمان برای Flash حدود ۱۳۸ میکروثانیه است.
گام بعدی شما
- اگر از مدل‌های سفارشی در PyTorch استفاده می‌کنید، حتماً خروجی
torch.profilerرا بررسی کنید تا مطمئن شوید بک‌اند 'math' به‌طور ناخواسته فعال نشده است. - برای کاهش تأخیر در محیط Production، اولویت خود را به ترتیب FlashAttention-2 و سپس cuDNN قرار دهید.
- در پیاده‌سازی‌های دستی، همیشه از نسخه‌های in-place مانند
masked_fill_استفاده کنید تا کپی‌های پنهان حافظه حذف شوند.
اما داستان سخت‌افزاری این تحول حتی شگفت‌انگیزتر است — به تحلیل ما درباره‌ی تراشه‌های Blackwell و مدیریت حافظه در نسل جدید مراجعه کنید.
گفتگو