FlashAttention-2 (FA)

From Systems analysis Wiki
Jump to navigation Jump to search

FlashAttention-2 — الگوریتمی پیشرفته است که برای محاسبه مکانیزم توجه (attention) در مدل‌های زبانی بزرگ (LLM) طراحی شده است. این الگوریتم توسط تری دائو (Tri Dao) و پژوهشگران دانشگاه استنفورد توسعه یافته و در ژوئیه ۲۰۲۳ معرفی شد[1]. هدف اصلی آن — تسریع قابل توجه آموزش و استنتاج (inference) مدل‌های transformer از طریق استفاده کارآمدتر از منابع سخت‌افزاری GPU، با حفظ کامل یکسانی محاسبات با مکانیزم توجه استاندارد، یعنی بدون افت دقت، می‌باشد.

FlashAttention-2 ادامه منطقی الگوریتم FlashAttention است که در سال ۲۰۲۲ توسط همان تیم معرفی شد. نسخه جدید مشکل بارگذاری ناقص GPU که در نسخه پیشین مشاهده می‌شد را برطرف کرده و به افزایش سرعتی تقریباً دو برابری نسبت به نسخه اول دست می‌یابد.

پیش‌زمینه: مشکل توجه در transformer‌ها

مکانیزم self-attention استاندارد هنگام کار با دنباله‌های متنی طولانی در transformer‌ها گلوگاه محسوب می‌شود. پیچیدگی محاسباتی و مصرف حافظه آن به صورت درجه دوم (O(N²)) بر حسب طول دنباله (N) رشد می‌کند، که محدودیت‌های جدی بر حداکثر طول زمینه و مقیاس‌پذیری LLM‌ها اعمال می‌کند[1].

برای حل این مشکل در سال ۲۰۲۲ الگوریتم FlashAttention معرفی شد[2]. ایده‌های کلیدی آن:

  • در نظر گرفتن سلسله‌مراتب حافظه GPU (IO-awareness): الگوریتم عملیات پرهزینه خواندن/نوشتن بین حافظه کُند GPU (HBM) و حافظه استاتیک سریع (SRAM) روی تراشه را به حداقل می‌رساند.
  • پردازش بلوکی (tiling): محاسبات به بلوک‌های کوچک (tile) تقسیم می‌شوند که در SRAM سریع پردازش می‌شوند و این امکان را می‌دهد که از مادی‌سازی ماتریس کامل توجه در حافظه اجتناب شود.

این رویکرد به دستیابی به رشد خطی مصرف حافظه (O(N)) و تسریع ۲ تا ۴ برابری نسبت به پیاده‌سازی‌های استاندارد منجر شد[2]. FlashAttention به طور گسترده‌ای به کار گرفته شد و به ظهور مدل‌هایی با زمینه به طور قابل توجهی افزایش‌یافته کمک کرد، به عنوان مثال از ۲ تا ۴ هزار token (GPT-3) تا ۱۲۸ هزار (GPT-4) و بیشتر[3]. در مدل Falcon-40B استفاده از FlashAttention سرعت استنتاج را ۳ برابر و عملکرد کلی تولید را ۵ برابر نسبت به GPT-3 افزایش داد[4].

توسعه و اهداف FlashAttention-2

با وجود موفقیت، نسخه اول FlashAttention از منابع محاسباتی GPU به طور کامل استفاده نمی‌کرد. در کارت‌های گرافیک NVIDIA A100 عملکرد تنها به ۲۵ تا ۴۰ درصد از حداکثر نظری (FLOPs/s) می‌رسید[1]. دلیل اصلی بارگذاری غیربهینه پردازنده‌های جریانی (Streaming Multiprocessors) و عملیات اضافی با حافظه مشترک بود[5].

هدف FlashAttention-2 تسریع بیشتر محاسبات از طریق موازی‌سازی کارآمدتر و به حداقل رساندن عملیات کمکی بود. الگوریتم به طور کامل با استفاده از عناصر اولیه سطح پایین کتابخانه NVIDIA CUTLASS 3.x برای دستیابی به حداکثر عملکرد بازنویسی شد[6].

معماری فنی و اصول کار

FlashAttention-2 سه بهبود کلیدی برای افزایش موازی‌سازی و کارایی معرفی می‌کند[1]:

۱. به حداقل رساندن عملیات غیر-ماتریسی

الگوریتم تعداد عملیات کمکی اعشاری که ضرب ماتریسی نیستند (non-matmul FLOPs) را کاهش می‌دهد. از آنجا که هسته‌های تنسوری GPU دقیقاً برای عملیات ماتریسی (GEMM) بهینه شده‌اند و آن‌ها را تا ۱۶ برابر سریع‌تر انجام می‌دهند، این تغییر امکان می‌دهد بیشتر اوقات از پربازده‌ترین بلوک‌های GPU استفاده شود.

۲. موازی‌سازی بهبودیافته

در FlashAttention اصلی، کار روی یک «سر» توجه موازی‌سازی نمی‌شد که منجر به بیکاری در هنگام دنباله‌های طولانی و اندازه‌های کوچک batch می‌شد. FlashAttention-2 موازی‌سازی بین‌بلوکی را معرفی می‌کند: اکنون محاسبات برای یک سر توجه بین پردازنده‌های جریانی مختلف GPU توزیع می‌شود که بارگذاری آن‌ها را به طور قابل توجهی افزایش می‌دهد.

۳. تقسیم بهینه‌شده کار درون بلوک

در سطح یک بلوک محاسباتی، کار بین گروه‌های thread (warp) برای کاهش تبادل داده از طریق حافظه مشترک (shared memory) توزیع مجدد شد. این امر تعداد عملیات اضافی خواندن/نوشتن مورد نیاز برای نرمال‌سازی Softmax را کاهش می‌دهد.

عملکرد و کارایی

به لطف بهبودهای معماری، FlashAttention-2 افزایش عملکرد قابل توجهی را نشان می‌دهد:

  • تسریع دو برابری: الگوریتم تقریباً ۲ برابر سریع‌تر نسبت به نسخه اول FlashAttention کار می‌کند[1].
  • استفاده بالا از GPU: روی GPU NVIDIA A100 به ۵۰ تا ۷۳ درصد از حداکثر نظری پهنای باند (TFLOPs) دست می‌یابد که به کارایی عملیات بهینه‌شده ضرب ماتریسی (GEMM) نزدیک است[1].
  • سرعت محاسبات رکوردی:
    • روی GPU A100 سرعتی تا 225 TFLOP/s در چرخه آموزش سرتاسری مدلی از نوع GPT به دست می‌آید که معادل ۷۲ درصد استفاده از بلوک‌های محاسباتی است. برای مقایسه، توجه استاندارد در همان شرایط GPU را با کمتر از ۱۰۰ TFLOP/s بارگذاری می‌کرد[7].
    • روی GPU H100 عملکرد به 335 TFLOP/s می‌رسد[7].

چنین افزایش عملکردی به عنوان مثال امکان می‌دهد مدلی با پنجره زمینه 16k token در همان زمانی آموزش داده شود که پیشتر برای پنجره 8k token نیاز بود[5]. مهم است که الگوریتم دقیق و قطعی باقی می‌ماند، بنابراین کاربرد آن بر کیفیت پیش‌بینی‌های مدل تأثیری ندارد[8].

کاربرد و یکپارچگی در اکوسیستم

FlashAttention-2 به سرعت به ابزاری استاندارد در اکوسیستم LLM تبدیل شد. این الگوریتم در بسیاری از framework‌ها و کتابخانه‌های محبوب یکپارچه شده است:

  • PyTorch: پشتیبانی بومی.
  • Hugging Face Transformers: پشتیبانی با پارامتر `attn_implementation="flash_attention_2"` هنگام بارگذاری مدل فعال می‌شود[9]. با ده‌ها معماری سازگار است (GPT، Llama، Falcon، BERT و غیره)[10].
  • TensorRT-LLM، xFormers و Triton: الگوریتم برای این پلتفرم‌ها پیاده‌سازی شده که استفاده گسترده را تضمین می‌کند[7].

یکپارچگی امکان ترکیب آسان FlashAttention-2 با سایر روش‌های بهینه‌سازی مانند کوانتیزاسیون (GPTQ، QLoRA) و fine-tuning کارآمد (PEFT) را فراهم می‌کند[9].

مقایسه با نسخه‌های بعدی

FlashAttention-3

پژوهش‌ها در زمینه بهینه‌سازی توجه ادامه دارند. در ژوئیه ۲۰۲۴ تری دائو FlashAttention-3 را معرفی کرد که هدف آن بهره‌برداری از قابلیت‌های معماری GPU NVIDIA Hopper (H100/H200) است. نوآوری‌های کلیدی[3]:

  • پشتیبانی از FP8: از محاسبات اعشاری ۸ بیتی برای تسریع بیشتر استفاده می‌کند.
  • عملیات ناهمزمان: از قابلیت‌های ناهمزمان GPU به صورت کارآمدتری بهره می‌برد.

FlashAttention-3 تسریع ۱.۵ تا ۲ برابری نسبت به FlashAttention-2 روی GPU H100 ارائه می‌دهد و به عملکردی تا 740 TFLOP/s (۷۵ درصد از حداکثر نظری) می‌رسد[11].

منابع

  • Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
  • Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
  • Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608.
  • Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
  • Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
  • Chen, Y. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
  • Liu, Y. et al. (2024). FastAttention: Extending FlashAttention-2 to NPUs and Low-Resource GPUs. OpenReview: 76NYyOrnfk.
  • Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
  • Wang, G. et al. (2024). FlashMask: Efficient and Rich Mask Extension of FlashAttention. OpenReview: rog0J435OO.
  • Abbott, V.; Zardini, G. (2025). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. arXiv:2412.03317.

یادداشت‌ها

  1. ↑ 1.0 1.1 1.2 1.3 1.4 1.5 Дао, Три. «FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning». arXiv:2307.08691 [cs.LG], 17 июля 2023 г. [۱]
  2. ↑ 2.0 2.1 «Optimizing LLMs for Speed and Memory». Hugging Face Documentation. [۲]
  3. ↑ 3.0 3.1 Дао, Три. «FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision». Tri Dao's Blog. [۳]
  4. ↑ «FlashAttention vs FlashAttention-2 - an Analysis». E2E Networks Blog. [۴]
  5. ↑ 5.0 5.1 «FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning». OpenReview. [۵]
  6. ↑ «FlashAttention-2». Hazy Research, Stanford University. [۶]
  7. ↑ 7.0 7.1 7.2 Дао, Три. «FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning» (PDF). arXiv:2307.08691. [۷]
  8. ↑ Рашка, Себастьян. «Llama 2 and FlashAttention 2». Ahead of AI Magazine. [۸]
  9. ↑ 9.0 9.1 Белькада, Юнес. «Faster and more memory efficient models with Flash Attention 2!». LinkedIn. [۹]
  10. ↑ «GPU inference». Hugging Face Documentation. [۱۰]
  11. ↑ Дао, Три, и др. «FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision». arXiv:2407.08608 [cs.LG], 11 июля 2024 г. [۱۱]
Retrieved from "https://systems-analysis.info/int/index.php?title=FlashAttention-2_(FA)&oldid=2199"