FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao; Daniel Y. Fu; Stefano Ermon; Atri Rudra; Christopher Ré

الأوراق البحثية · arXiv 2205.14135

FlashAttention: انتباه تام سريع وعالي الكفاءة في استخدام الذاكرة مع مراعاة عمليات الإدخال والإخراج (IO-Awareness)

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré

تاريخ النشر 2022-05-27cs.LG arXiv PDF

اقرأ مع الترجمة والذكاء الاصطناعي

الملخص في ثلاثة أسطر

🎯 يجسد الانتباه التام القياسي مصفوفة انتباه بتعقيد $O (N^{2})$ في ذاكرة GPU HBM البطيئة، مما يجعلها محدودة بسرعة نقل الذاكرة ويصعب توسيع نطاقها للسياقات الطويلة.
🔧 تحسب FlashAttention الانتباه التام ضمن kernel موحد ومدمج في CUDA عبر تبليط المدخلات على ذاكرة SRAM السريعة الموجودة على الرقاقة، وتحديث softmax تراكمياً، وإعادة حساب الانتباه على الرقاقة أثناء التمرير العكسي (backward pass).
📈 تحقق FlashAttention تسريعاً كلياً في التدريب يصل إلى 3 أضعاف على GPT-2، وتدريباً أسرع بنسبة 15% مقارنة بالرقم القياسي لـ BERT-large في MLPerf 1.1، وتوسع نطاق السياق حتى 64K رمزاً.

الملخص

تتسم نماذج Transformers بالبطء والاستهلاك العالي للذاكرة في السلاسل الطويلة، نظراً لأن التعقيد الزمني والمكاني للانتباه الذاتي (self-attention) تربيعي بالنسبة لطول السلسلة. حاولت طرق الانتباه التقريبي معالجة هذه المشكلة عن طريق التضحية بجودة النموذج لتقليل التعقيد الحسابي، ولكنها غالباً لا تحقق تسريعاً زمنياً فعلياً على أرض الواقع (wall-clock speedup). ونرى أن هناك مبدأً مفقوداً يتمثل في جعل خوارزميات الانتباه واعية بعمليات الإدخال والإخراج (IO-aware) — أي مراعاة عمليات القراءة والكتابة بين مستويات الذاكرة المختلفة لوحدة معالجة الرسومات (GPU). نقترح خوارزمية FlashAttention، وهي خوارزمية انتباه تام واعية بعمليات الإدخال والإخراج تستخدم تقنية التبليط (tiling) لتقليل عدد عمليات القراءة/الكتابة بين الذاكرة ذات النطاق الترددي العالي (HBM) للـ GPU والذاكرة السريعة على الرقاقة (on-chip SRAM). نقوم بتحليل تعقيد الإدخال والإخراج لخوارزمية FlashAttention، مبيّنين أنها تتطلب وصولاً أقل إلى HBM مقارنة بالانتباه القياسي، وأنها مثالية لمجموعة واسعة من أحجام SRAM. نوسع أيضاً FlashAttention إلى الانتباه المتفرق القائم على الكتل (block-sparse attention)، مما ينتج خوارزمية انتباه تقريبي أسرع من أي طريقة انتباه تقريبي موجودة. تدرب FlashAttention نماذج Transformers بسرعة أكبر مقارنة بالنماذج المرجعية الحالية: تسريع فعلي شامل بنسبة 15% على نموذج BERT-large (طول السلسلة 512) مقارنة بالرقم القياسي لسرعة التدريب في MLPerf 1.1، وتسريع بمقدار 3 $\times$ على GPT-2 (طول السلسلة 1K)، وتسريع بمقدار 2.4 $\times$ على Long-Range Arena (طول السلسلة 1K-4K). تتيح كل من FlashAttention وblock-sparse FlashAttention سياقاً أطول في نماذج Transformers، مما يمنح نماذج ذات جودة أعلى (تحسن بمقدار 0.7 في perplexity لنموذج GPT-2 وارتفاع 6.4 نقطة في تصنيف المستندات الطويلة) وقدرات جديدة كلياً: تحقيق أول نماذج Transformers لأداء يتجاوز التخمين العشوائي في تحدي Path-X (طول السلسلة 16K، بدقة 61.4%) وPath-256 (طول السلسلة 64K، بدقة 63.1%).

الملخص الأصلي (بالإنجليزية)

Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware -- accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3 $\times$ speedup on GPT-2 (seq. length 1K), and 2.4 $\times$ speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).

دليل القراءة

المشكلة والدوافع

تمتلك نماذج Transformers تعقيداً زمنياً ومكانياً تربيعياً بالنسبة لطول السلسلة $N$ . على الرغم من أن طرق الانتباه التقريبي تقلل من عدد عمليات الفاصلة العائمة (FLOPs) نظرياً، إلا أنها غالباً ما تفشل في تحقيق تسريع فعلي على أرض الواقع لأن حسابات وحدات معالجة الرسومات الحديثة تواجه اختناقاً أساسياً في عمليات نقل الذاكرة (IO) وليس في العمليات الحسابية.

في الانتباه القياسي، تؤدي العمليات الحسابية $S = Q K^{⊤}$ و $P = softmax (S)$ و $O = PV$ إلى كتابة مصفوفات وسيطة بحجم $N \times N$ في ذاكرة HBM البطيئة وقراءتها مجدداً عدة مرات. هذا يتسبب في وصول لذاكرة HBM بتعقيد $Θ (N d + N^{2})$ واستهلاك ذاكرة بمقدار $O (N^{2})$ ، مما يحد من نوافذ السياق العملية.

الأفكار الرئيسية

مراعاة عمليات الإدخال والإخراج (IO-Awareness): مراعاة التسلسل الهرمي لذاكرة GPU بشكل صريح: ذاكرة SRAM السريعة والصغيرة على الرقاقة (مثل 192 كيلوبايت لكل معالج تيار متعدد SM بنطاق ترددي يقارب 19 تيرابايت/ثانية في بطاقة A100) مقابل ذاكرة HBM الكبيرة والبطيئة (40–80 جيجابايت بنطاق ترددي 1.5–2.0 تيرابايت/ثانية).
التبليط (حساب Softmax المجزأ في كتل): تقسيم المدخلات $Q, K, V$ إلى كتل تتسع بالكامل داخل ذاكرة SRAM. ومن خلال تتبع إحصائيات القيمة العظمى للصفوف وعوامل التسوية الجارية لـ softmax، يتم حساب اختزال softmax بشكل تراكمي عبر الكتل دون تجسيد المصفوفة الكاملة بحجم $N \times N$ في ذاكرة HBM.
إعادة الحساب في التمرير العكسي: تجنب حفظ مصفوفة الانتباه الوسيطة بحجم $N \times N$ من أجل التمرير العكسي (backpropagation). بدلاً من ذلك، يتم حفظ المخرجات $O$ وإحصائيات softmax $(ℓ, m)$ فقط، مع إعادة حساب $S$ و $P$ على الرقاقة كتلة تلو الأخرى داخل SRAM أثناء التمرير العكسي.
دمج الـ Kernel (Kernel Fusion): دمج عمليات ضرب المصفوفات، وحجب القيم (masking)، وsoftmax، وdropout في نواة CUDA واحدة للقضاء على رحلات الذهاب والإياب الزائدة إلى HBM.

الطريقة (مع المعادلات الرئيسية)

بفرض وجود $Q, K, V \in R^{N \times d}$ ، تحدد FlashAttention أحجام الكتل بـ $B_{c} = ⌈ M / (4 d)⌉$ و $B_{r} = min (⌈ M / (4 d)⌉, d)$ ، حيث يمثل $M$ سعة ذاكرة SRAM.

التتبع التراكمي لـ Softmax

بالنسبة لمتجه مجزأ إلى متجهات فرعية $x^{(1)}$ و $x^{(2)}$ ، يتم تحديث الإحصائيات الجارية عبر: $m (x) = max (m (x^{(1)}), m (x^{(2)})),$ $ℓ (x) = e^{m (x^{(1)}) - m (x)} ℓ (x^{(1)}) + e^{m (x^{(2)}) - m (x)} ℓ (x^{(2)}) .$

لكل كتلة $j$ من المفاتيح/القيم والكتلة $i$ من الاستعلامات:

حساب الدرجات غير المعيرة داخل SRAM: $S_{ij} = Q_{i} K_{j}^{⊤} \in R^{B_{r} \times B_{c}}$ .
حساب الإحصائيات المحلية للكتلة: $\tilde{m}_{ij} = rowmax (S_{ij})$ ، و $\tilde{P}_{ij} = exp (S_{ij} - \tilde{m}_{ij})$ ، و $\tilde{ℓ}_{ij} = rowsum (\tilde{P}_{ij})$ .
تحديث عوامل التسوية العامة الجارية: $m_{i}^{new} = max (m_{i}, \tilde{m}_{ij}), ℓ_{i}^{new} = e^{m_{i} - m_{i}^{new}} ℓ_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{ℓ}_{ij} .$
تحديث كتلة المخرجات الجارية $O_{i} \in R^{B_{r} \times d}$ : $O_{i} \leftarrow diag (ℓ_{i}^{new})^{- 1} (diag (ℓ_{i}) e^{m_{i} - m_{i}^{new}} O_{i} + e^{\tilde{m}_{ij} - m_{i}^{new}} \tilde{P}_{ij} V_{j}) .$

تحليل التعقيد

تعقيد الإدخال والإخراج (IO Complexity): يتطلب الانتباه القياسي عمليات وصول إلى HBM بتعقيد $Θ (N d + N^{2})$ . بينما تتطلب FlashAttention عمليات وصول بتعقيد $Θ (N^{2} d^{2} M^{- 1})$ . يثبت المؤلفون أنه لا يمكن لأي خوارزمية انتباه تام أن تحقق خطياً مقارباً $o (N^{2} d^{2} M^{- 1})$ من عمليات الوصول لجميع قيم $M \in [d, N d]$ .
العمليات الحسابية والذاكرة: تنفذ FlashAttention عمليات FLOPs بتعقيد $O (N^{2} d)$ (مع عمليات حسابية إضافية طفيفة في التمرير العكسي نتيجة إعادة الحساب) وتتطلب ذاكرة مساعدة بتعقيد $O (N)$ فقط بخلاف المدخلات والمخرجات.

الامتداد: Block-Sparse FlashAttention

يؤدي قصر الحسابات على الكتل غير الصفرية وفقاً لقناع تفرق مسبق $M \in {0, 1}^{N / B_{r} \times N / B_{c}}$ إلى تعقيد وصول لذاكرة HBM مقداره $Θ (N d + N^{2} d^{2} M^{- 1} s)$ ، حيث يمثل $s$ نسبة الكتل غير الصفرية.

النتائج (مع أهم الأرقام)

تسريع التدريب:
- BERT-large (طول السلسلة 512): تفوق على الرقم القياسي لسرعة التدريب في MLPerf 1.1 بنسبة 15% (17.4 دقيقة مقابل 20.0 دقيقة على 8 بطاقات A100).
- GPT-2 (طول السلسلة 1K): حقق تسريعاً يصل إلى 3.5 أضعاف مقارنة بـ HuggingFace وما يصل إلى 1.8 أضعاف مقارنة بـ Megatron-LM.
- Long-Range Arena (LRA): حقق تسريعاً بمقدار 2.4 أضعاف مقارنة بالانتباه القياسي، بينما حققت نسخة Block-sparse FlashAttention تسريعاً بمقدار 2.8 أضعاف.
إمكانات السياق الأطول:
- GPT-2 بسياق 4K: يتدرب أسرع من Megatron-LM بسياق 1K مع تحسين perplexity للتحقق بمقدار 0.7 نقطة (17.5 مقابل 18.2).
- تصنيف المستندات الطويلة: أدى توسيع السياق من 512 إلى 16K إلى تحسن قدره 4.3 في مقياس micro- $F_{1}$ على MIMIC-III (من 52.8 إلى 57.1) وما يصل إلى 8.5 على ECtHR (من 72.2 إلى 80.7 عند سياق 8K).
- Path-X وPath-256: تُعد FlashAttention أول بنية Transformer تنجح في حل اختبار Path-X (طول السلسلة 16K) بدقة 61.4%. وحققت Block-sparse FlashAttention دقة 63.1% على Path-256 (طول السلسلة 64K).

القيود والأسئلة المفتوحة

أعباء هندسة البرمجيات في CUDA: يتطلب تنفيذ نوى (kernels) واعية بالإدخال والإخراج برمجة مخصصة منخفضة المستوى بلغة CUDA بدلاً من تجريدات PyTorch القياسية أو المترجمات (compilers)، مما يجعل قابلية النقل عبر معماريات العتاد المختلفة أمراً صعباً.
الانتباه عبر وحدات معالجة رسومات متعددة: لا يمتد جدول عمليات الإدخال والإخراج الأمثل للعقدة الواحدة تلقائياً إلى معمارية الانتباه الموزع عبر عدة GPUs والروابط البينية (مثل NVLink وInfiniBand).
الحد الأدنى للانتباه التام: على الرغم من أن الخوارزمية مثالية بالنسبة للقيم العامة لـ $M$ ، فإن إثبات حدود دنيا دقيقة للتعقيد المتغير بدلالة $M$ يظل مسألة بحثية مفتوحة.

المصطلحات الرئيسية

IO-awareness (مراعاة عمليات الإدخال والإخراج): المبدأ الخوارزمي لتصميم العمليات بحيث تراعي عمليات القراءة والكتابة بين المستويات المختلفة لذاكرة العتاد (مثل ذاكرة SRAM السريعة على الرقاقة مقابل ذاكرة HBM الأبطأ).
High Bandwidth Memory (HBM) (الذاكرة ذات النطاق الترددي العالي): الذاكرة الأساسية الكبيرة لوحدة معالجة الرسومات (GPU) التي توفر سعة إجمالية عالية ولكن بنطاق ترددي أقل بكثير مقارنة بذاكرة التخزين المؤقت المدمجة بالرقاقة.
SRAM: ذاكرة وسيطة سريعة ومدمجة على الرقاقة في وحدات معالجة الرسومات الحديثة تقع داخل كل معالج تدفق متعدد (SM).
Tiling (التبليط): إعادة هيكلة العملية الحسابية عبر تقسيم مصفوفات المدخلات إلى كتل أصغر تتسع داخل ذاكرة SRAM وتجميع النتائج عبر مراحل متتابعة.
Selective gradient checkpointing (نقاط التحقق الانتقائية للتدرجات): تقنية تتجنب حفظ مصفوفة الانتباه الوسيطة الكبيرة أثناء التمرير الأمامي، وذلك عبر إعادة حسابها على الرقاقة أثناء التمرير العكسي باستخدام إحصائيات softmax المخزنة مؤقتاً.
Kernel fusion (دمج العمليات في الـ Kernel): دمج عمليات متباينة متعددة (ضرب المصفوفات، والحجب، وsoftmax، وdropout) في تنفيذ واحد لـ kernel على الـ GPU للقضاء على عمليات القراءة والكتابة المتكررة إلى HBM.
Block-sparse FlashAttention: امتداد لخوارزمية FlashAttention يحسب فقط كتل التبليط غير الصفرية المحددة بواسطة قناع التفرق الكتلوي، مما يقلل تعقيد الإدخال والإخراج بالتناسب مع نسبة التفرق.
Path-X / Path-256: اختبارات معيارية للاستدلال البصري ذي السلاسل الطويلة ضمن Long-Range Arena (بأطوال سلاسل 16K و64K) فشلت فيها نماذج Transformers السابقة في تحقيق دقة تتجاوز التخمين العشوائي.

أسئلة شائعة

هل تعد خوارزمية FlashAttention تقريباً لمصفوفة الانتباه؟

لا. تحسب FlashAttention انتباهاً ذاتياً تاماً ينتج عنه مخرجات رياضية مطابقة تماماً للانتباه القياسي حتى حدود دقة الفاصلة العائمة العددية.

لماذا تكون FlashAttention أسرع رغم أنها تحسب عمليات FLOPs أكثر في التمرير العكسي؟

لأن الانتباه القياسي محدود بسرعة الذاكرة؛ إذ يستهلك نقل مصفوفة $N \times N$ بين HBM وSRAM وقتاً زمنياً فعلياً أطول بكثير من إجراء عمليات حسابية إضافية داخل ذاكرة SRAM السريعة.

ما هي الذاكرة المساعدة التي تخزنها FlashAttention من أجل التمرير العكسي؟

تخزن FlashAttention مخرجات التمرير الأمامي $O$ فقط، وعوامل تسوية softmax المتمثلة في $ℓ$ و $m$ (والتي تتطلب ذاكرة إضافية بتعقيد $O (N)$ )، وحالة المولد شبه العشوائي لإعادة إنتاج أقنعة dropout على الرقاقة.

كيف تقارن FlashAttention بالانتباه القياسي من حيث تعقيد الإدخال والإخراج (IO)؟

يتطلب الانتباه القياسي وصولاً لذاكرة HBM مقداره $Θ (N d + N^{2})$ ، بينما تتطلب FlashAttention وصولاً مقداره $Θ (N^{2} d^{2} M^{- 1})$ ، حيث $M$ يمثل حجم SRAM و $d$ يمثل بُعد الرأس.

كتب الذكاء الاصطناعي الشرح في هذه الصفحة استنادًا إلى الورقة، وقد يحتوي على أخطاء. راجع الورقة للتفاصيل.

اقرأ مع الترجمة والذكاء الاصطناعي