FlashAttention-4: كيف سرعت Together AI خوارزمية الانتباه على معالج Blackwell

FlashAttention-4 هي إعادة تصميم كاملة لخوارزمية الانتباه خصيصاً للمعالجات الحديثة، حيث تنمو نوى الموتر في الأداء بسرعة أكبر بكثير من الذاكرة والموارد الأخرى.

لماذا لا يعمل النهج القياسي بعد الآن

للوهلة الأولى، يتحكم في أداء الانتباه سرعة عمليتي ضرب المصفوفات: S = Q × K^T و O = P × V. لكن تحليل Blackwell B200 يكشف عن شيء غير متوقع: الاختناق ليس في نوى الموتر، بل في وحدات الدوال المتخصصة (SFU) لحساب الدالة الأسية في softmax (المرحلة الأمامية) وفي حركة الذاكرة المشتركة (المرحلة العكسية). من Hopper (H100) إلى Blackwell (B200)، ارتفعت أداء نوى الموتر BF16 من 1 إلى 2.25 بيتافلوب/الثانية، لكن عدد وحدات SFU وعرض نطاق الذاكرة المشتركة ظلا دون تغيير. هذا عدم التناسق يكسر التحسين القياسي — لا يمكن افتراض أن النوى تحدد الأداء بالكامل. في الواقع، يتداخلان مع بعضهما البعض.

كيف حلت Together AI المشكلة

اقترحت مجموعة البحث بالتعاون مع مهندسي NVIDIA و Meta و Princeton ثلاث أفكار رئيسية:

• خط أنابيب جديد — خطوط أنابيب برمجية تتداخل بأقصى قدر مع عمل نوى الموتر ووحدات SFU والذاكرة بدون توقف.
• محاكاة الدالة الأسية عبر كثيرات الحدود — بدلاً من كتلة SFU البطيئة، تحسب المرحلة الأمامية الدالة الأسية على وحدات FMA سريعة (عمليات الضرب والجمع المدمجة).
• TMEM و 2-CTA MMA — استخدام ذاكرة الموتر الجديدة (256 كيلوبايت لكل وحدة معالجة) بالإضافة إلى طريقة حيث تعمل كتلتا خيط على عملية مصفوفية واحدة، مما يقلل حركة الذاكرة المشتركة.

قدرات Blackwell التي جعلت هذا ممكناً

أضافت Blackwell عدة ميزات خصيصاً لمثل هذه التحسينات:

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

نوى الموتر غير المتزامنة من الجيل الخامس — كل نواة تبدأ بخيط واحد وتجمع النتائج في TMEM. أقصى حجم مربع لـ BF16 هو 128×256×16 (أكبر تقريباً مرتين من Hopper)، مما يسمح بخط أنابيب أعمق دون امتلاء السجلات.

2-CTA MMA — طريقة جديدة حيث تعمل كتلتا خيط معاً على عملية مصفوفية واحدة. هذا يقلل الحمل على الذاكرة المشتركة بمقدار النصف ويقلل عدد العمليات الذرية.

الأرقام والنتائج

تصل FlashAttention-4 على Blackwell B200 مع BF16 إلى 1605 تيرافلوب/الثانية (بنسبة استخدام 71%). هذا أسرع بـ 1.3 مرة من cuDNN 9.13 وأسرع بـ 2.7 مرة من Triton. للمقارنة: 1605 تيرافلوب/الثانية — وهي تقريباً نصف الأداء القصوى لـ Blackwell، لكنها تُستخرج من نواة الانتباه المعقدة.

ماذا يعني هذا

تُظهر FlashAttention-4 كيف يجب العمل في عصر التوسع غير المتناسق لمعالجات GPU — لا تُحكم براغي الخوارزميات القديمة، بل أعد تصميمها مع الإمكانيات الجديدة للأجهزة. تبدو النتيجة كإنجاز، لكنها في الواقع مجرد بداية مواءمة مع الواقع الجديد للأجهزة.

*تُعتبر Meta منظمة متطرفة وممنوعة في روسيا الاتحادية.