FlashAttention-2: انتباه أسرع بتوازٍ أفضل وتوزيع عمل أذكى

المعالجات الرسومية الحديثة ليست مجرد آلات حساب سريعة — بل منظومات معقّدة ذات هرميّة ذاكرة متعدّدة المستويات. في المستوى الأبعد تقع ذاكرة HBM: سعتها كبيرة (40‑80 غيغابايت في A100) لكنها أبطأ نسبياً (1.5‑2.0 تيرابايت/ثانية). أما على الشريحة نفسها فتوجد ذاكرة SRAM: حجمها ضئيل (192 كيلوبايت لكل وحدة معالجة متدفّقة) لكنها فائقة السرعة (نحو 19 تيرابايت/ثانية). فرق السرعة بينهما يبلغ 10 أضعاف تقريباً.

الفكرة المحورية في FlashAttention وFlashAttention-2 هي أن حساب الانتباه مقيّد أساساً بنقل البيانات لا بالحساب ذاته. المعالج يُمضي معظم وقته في نقل البيانات بين HBM وSRAM، لا في إجراء العمليات الحسابية فعلياً. التنفيذ التقليدي يبني مصفوفة الانتباه الكاملة $N \times N$ ومخرج softmax في ذاكرة HBM — أي أنه يقرأ ويكتب $O(N^2)$ عنصر لتسلسل بطول $N$. وحين يتجاوز $N$ قيمة 4096، يصبح نقل البيانات هو العائق الحقيقي لا الحساب.

المعالجات الرسومية تمتلك أيضاً نوى مُوتِّرات مُصمَّمة خصيصاً لضرب المصفوفات. على A100 تُنجز هذه النوى 312 تيرا عملية/ثانية بنصف الدقة (FP16/BF16)، بينما الحساب العام بدقة FP32 لا يتجاوز 19.5 تيرا عملية/ثانية. بمعنى آخر، كل عملية غير مصفوفية تُكلّف عملياً 16 ضعف تكلفة عملية ضرب مصفوفات. فلسفة FlashAttention-2 مبنيّة على هذا التباين: اجعل المعالج يقضي أكبر وقت ممكن في ضرب المصفوفات، وقلّل كل ما عدا ذلك إلى الحدّ الأدنى.

قبل الدخول في التحسينات، لنُراجع بإيجاز ما يفعله FlashAttention. الانتباه بصورته التقليدية يحسب $\text{S} = QK^\top$ ثم يُطبّق softmax على كل صف ليحصل على $P$، ثم يحسب $O = PV$. المشكلة أن هذا يتطلّب بناء مصفوفتَي $N \times N$ كاملتين (S وP) في ذاكرة HBM.

FlashAttention يتفادى ذلك عبر أسلوب التبليط: يُقسّم Q وK وV إلى كتل صغيرة تتّسع في SRAM. لكل زوج كتل، يحسب درجات الانتباه محلياً ويُطبّق softmax ويضرب في كتلة V المقابلة ثم يُراكم النتيجة. التحدّي هنا أن softmax عملية شاملة — تحتاج القيمة القصوى والمجموع على الصف بأكمله. تقنية الـ softmax المتدفّق تحلّ ذلك عبر الاحتفاظ بإحصاءات جارية: قيمة قصوى $m$ ومجموع أُسّي $\ell$ يُحدَّثان مع كل كتلة جديدة. عند وصول كتلة جديدة تُحدَّث الإحصاءات ويُعاد ضبط المخرج الجزئي، فتبقى النتيجة النهائية مطابقة تماماً — بدون أي تقريب.

بهذا الأسلوب تنخفض الذاكرة من $O(N^2)$ إلى $O(N)$ ويتحقّق تسريع بمقدار 2‑4× بفضل تقليل القراءة والكتابة من وإلى HBM. لكنّ هناك أداءً مُهدَراً لا يزال يمكن استرجاعه.

التحسين الأول يستهدف خطوة التطبيع في الـ softmax المتدفّق. في FlashAttention الأصلي، بعد معالجة كل كتلة K-V جديدة، يُطبَّع كلا حدَّي تحديث المخرج بالقسمة على $\text{diag}(\ell^{(j)})^{-1}$:

$O^{(j)} = \text{diag}(\ell^{(j-1)} / \ell^{(j)}) \cdot O^{(j-1)} + \text{diag}(\ell^{(j)})^{-1} \cdot e^{S^{(j)} - m^{(j)}} V^{(j)}$

الملاحظة المفتاحية في FlashAttention-2 هي أن هذا التطبيع يمكن تأجيله. بدلاً من القسمة على $\ell$ في كل خطوة، يحتفظ الخوارزم بمخرج جارٍ غير مُطبَّع $\tilde{O}$ ولا يقسم إلا مرة واحدة في النهاية:

$\tilde{O}^{(j)} = \text{diag}(e^{m^{(j-1)} - m^{(j)}}) \cdot \tilde{O}^{(j-1)} + e^{S^{(j)} - m^{(j)}} V^{(j)}$

بهذا تُحذف عملية قسمة واحدة من كل كتلة — وهي عملية غير مصفوفية مُكلفة. علاوة على ذلك، بدلاً من تخزين القيمة القصوى $m$ والمجموع $\ell$ كلٍّ على حدة لـالتمرير العكسي، يُخزّن FlashAttention-2 حاصل دمجهما في قيمة واحدة هي الـ logsumexp: $L = m + \log(\ell)$. هذا يُنصّف الحِمل الحسابي الإضافي ويُبسّط التمرير العكسي.

في FlashAttention الأصلي، يتوزّع العمل على بُعدَين فقط: حجم الدُّفعة وعدد رؤوس الانتباه. كل كتلة خيوط تتولّى رأس انتباه واحداً من تسلسل واحد. معالج A100 يضمّ 108 وحدة معالجة متدفّقة (SMs)، وتحتاج إلى نحو 80 كتلة خيوط على الأقل لإشغاله فعلياً. فلو كان حجم الدُّفعة 1 وعدد الرؤوس 32، لن تحصل إلا على 32 كتلة — ما يترك وحدات معالجة كثيرة خاملة.

FlashAttention-2 يُضيف بُعداً ثالثاً للتوازي: طول التسلسل. الحلقة الخارجية تمرّ على كتل الصفوف من Q، وكل تكرار مستقلّ تماماً عن البقية (يقرأ K وV بالكامل لكنه يكتب في شريحته الخاصة من المخرج). هذا النوع من التوازي مثالي لأنه لا يتطلّب أي تنسيق بين الكتل.

بعد هذا التعديل يصبح العدد الكلّي لكتل الخيوط الدُّفعة × الرؤوس × (N / حجم_الكتلة)، وهو ما قد يبلغ الآلاف في التسلسلات الطويلة. حتى تسلسل طويل واحد برأس انتباه واحد يكفي الآن لإشغال المعالج بالكامل. وهذا هو السيناريو الأكثر أهمية عملياً — تدريب السياقات الطويلة مع نماذج كبيرة ودُفعات صغيرة.

في التمرير العكسي يعمل التوازي بأسلوب مشابه لكن مع فارق مهم. الحلقة الخارجية هنا تمرّ على كتل الأعمدة من K وV (لا كتل الصفوف من Q). كل كتلة عمود تُراكم dK وdV باستقلالية تامة. الاستثناء هو dQ: تحديثه مُشترك بين كتل الأعمدة، إذ قد تحتاج عدة كتل خيوط إلى الكتابة في الشريحة نفسها من dQ. الحل هو استخدام الجمع الذرّي (atomic addition)، وهو آلية تزامن خفيفة تسمح بالتحديث المتزامن دون الحاجة لحواجز تزامن كاملة.

حتى داخل كتلة خيوط واحدة، هناك عدة أشرطة (كل شريط مجموعة من 32 خيطاً) تحتاج إلى تقاسم العمل فيما بينها. الطريقة التي تتقاسم بها هذا العمل تُحدّد كمّ التواصل المطلوب عبر الذاكرة المشتركة — وهذا التواصل مُكلف.

في FlashAttention كانت الاستراتيجية هي تقسيم المفاتيح (split-K): تُوزَّع K وV على 4 أشرطة، بينما تتشارك كلها في Q. كل شريط يحسب جزءاً من $QK^\top$ ثم يضربه في جزئه من V. المشكلة أن النتائج الجزئية تحتاج دمجاً — كل شريط يكتب نتيجته في الذاكرة المشتركة، ثم تتوقّف جميع الأشرطة للتزامن، ثم يقرأ شريط واحد الأجزاء ويجمعها. دورة الكتابة-التزامن-القراءة هذه هي عنق الزجاجة.

FlashAttention-2 يقلب المعادلة إلى تقسيم الاستعلامات (split-Q): تُوزَّع Q على الأشرطة بينما تتشارك كلها في K وV. الآن كل شريط يعمل على صفوفه الخاصة من Q باستخدام K وV كاملتين، ويُنتج شريحته من المخرج مباشرةً. لا حاجة لأي تواصل بين الأشرطة. الصورة الذهنية: أعطِ كل عامل قائمة أسئلته الخاصة (Q) ودعهم جميعاً يقرأون من المرجع نفسه (K وV) — كل عامل يكتب إجاباته وحده دون الحاجة لمساعدة أحد.

في نمذجة اللغة الارتدادية، يضمن القناع السببي أن كل رمز لا ينتبه إلا للرموز التي سبقته. على مستوى المصفوفة، كل عنصر فوق القطر يُصفَّر (يُعيَّن بقيمة $-\infty$). وبما أن FlashAttention يعمل بالكتل أصلاً، يمكنه تخطّي أي كتلة تقع بالكامل فوق القطر — أي نحو نصف الكتل في التسلسلات الطويلة. هذا التخطّي وحده يمنح تسريعاً بمقدار 1.7‑1.8×.

FlashAttention-2 يُحسّن الأمر أكثر: بالنسبة للكتل التي يمرّ القطر من خلالها، لا يُطبَّق القناع إلا داخل تلك الكتلة الواحدة التي تقع على حافة القطر. أما الكتل الواقعة بالكامل تحت القطر فلا تحتاج لأي تقنيع. النتيجة أن عبء القناع ينحصر في كتلة واحدة لكل صف بدلاً من تطبيقه عنصراً عنصراً على أجزاء كبيرة من المصفوفة.

حين تتضافر التحسينات الثلاثة، يكون الأثر كبيراً. على معالج A100 بسعة 80 غيغابايت، يصل FlashAttention-2 إلى 230 تيرا عملية/ثانية في التمرير الأمامي (73% من الحدّ النظري) و196 تيرا عملية/ثانية في التمرير العكسي (63%). في المقابل، FlashAttention الأصلي لا يتجاوز 124 أمامياً و113 عكسياً. أي أن التسريع يبلغ نحو ضعفين بشكل شبه ثابت.

وحين نقارن بالانتباه التقليدي في PyTorch، تتّسع الفجوة أكثر: FlashAttention-2 أسرع بـ 3‑10 أضعاف بحسب طول التسلسل وبُعد الرأس، مع استهلاك ذاكرة $O(N)$ بدلاً من $O(N^2)$.

على مستوى التدريب الكامل، يصل FlashAttention-2 إلى 225 تيرا عملية/ثانية لكل معالج A100 عند تدريب نماذج بأسلوب GPT بـ 2.7 مليار معامل وسياق بطول 8 آلاف رمز. هذا يُمثّل 72% من الإنتاجية النظرية للنموذج — وهي نسبة قريبة جداً من الحدّ العتادي. للمقارنة، التدريب نفسه بدون FlashAttention لا يتجاوز 80 تيرا عملية/ثانية، أي أبطأ بمقدار 2.8×. وكلّما طال التسلسل زاد المكسب: عند سياق 8 آلاف يبلغ التسريع 2.8× مقارنة بالتنفيذ التقليدي، بينما عند سياق 2 آلاف يبلغ 1.4× فقط.

نماذج اللغة الكبيرة الحديثة تعتمد بشكل متزايد على الانتباه متعدّد الاستعلامات (MQA) والانتباه المُجمَّع (GQA) لتقليص حجم ذاكرة KV المؤقتة أثناء الاستدلال. في MQA تتشارك جميع رؤوس الاستعلام في رأس مفاتيح-قيم واحد، وفي GQA تتشارك مجموعات من رؤوس الاستعلام في رأس واحد. كلا الأسلوبين يُقلّل حجم البيانات المنقولة أثناء التوليد.

FlashAttention-2 يدعم هذه الأنماط بشكل طبيعي: بدلاً من نسخ رؤوس K وV فعلياً لمطابقة عدد رؤوس Q، يستخدم مؤشرات فهرسة ذكية تجعل عدة مجموعات من Q تقرأ من بيانات K-V نفسها دون تكرارها. في التمرير العكسي تُجمَع التدرّجات dK وdV عبر رؤوس الاستعلام المتشاركة. أهمية هذا الدعم أنه يعني أن FlashAttention-2 لا يُسرّع التدريب فحسب، بل يُسرّع أيضاً بنيات الاستدلال المُستخدمة فعلياً في الإنتاج.

لنجمع كل ما سبق في صورة الخوارزم الكامل لـالتمرير الأمامي. الحلقة الخارجية تمرّ على كتل صفوف Q، وكل كتلة تُسنَد إلى كتلة خيوط مستقلة للتوازي. الحلقة الداخلية تمرّ على كتل أعمدة K وV. في كل خطوة يحسب الخوارزم درجات الانتباه المحلية ويُحدّث إحصاءات softmax الجارية ويُراكم المخرج دون تطبيع. التطبيع لا يحدث إلا مرة واحدة في النهاية.