FlashAttention: توجه دقیق سریع و کارآمد با IO-Awareness

FlashAttention یک الگوریتم نوآورانه است که برای رسیدگی به تنگناهای محاسباتی و حافظه مکانیزم توجه در مدل‌های ترانسفورمر، به ویژه برای دنباله‌های طولانی طراحی شده است. با بهینه‌سازی الگوهای دسترسی به حافظه در GPUها، سرعت و کارایی حافظه قابل توجهی را به دست می‌آورد و آن را به ابزاری اساسی برای مقیاس‌بندی ترانسفورمرها به طول‌های زمینه طولانی‌تر تبدیل می‌کند.

شکل 1: سمت چپ: FlashAttention از کاشی کاری برای جلوگیری از ماتریس توجه بزرگ × × (جعبه نقطه‌دار) در GPU HBM (نسبتاً) کند استفاده می‌کند. در حلقه بیرونی (فلش‌های قرمز)، FlashAttention از میان بلوک‌های ماتریس‌های K و V حلقه می‌زند و آن‌ها را برای SRAM سریع روی تراشه بارگذاری می‌کند. در هر بلوک، FlashAttention روی بلوک‌های ماتریس Q (فلش‌های آبی) حلقه می‌زند، آنها را در SRAM بارگذاری می‌کند و خروجی محاسبه توجه را به HBM باز می‌نویسد. راست: افزایش سرعت اجرای PyTorch از توجه در GPT-2. FlashAttention ماتریس توجه بزرگ