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

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

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

يبدو للوهلة الأولى أن أداء الانتباه (attention) تحكمه سرعة عمليتي ضرب مصفوفات: S = Q × K^T وO = P × V. لكن تحليل Blackwell B200 يكشف أمرًا غير متوقع: فعنق الزجاجة لا يقع في النوى التنسورية، بل في وحدات الدوال الخاصة (SFU) التي تحسب الدالة الأسية في softmax (في المرور الأمامي forward pass)، وفي حركة البيانات عبر الذاكرة المشتركة shared memory (في المرور العكسي backward pass).

من Hopper (H100) إلى Blackwell (B200) ارتفع أداء النوى التنسورية بصيغة BF16 من 1 إلى 2.25 بيتافلوب، بينما بقي عدد وحدات SFU وعرض نطاق الذاكرة المشتركة على حالهما. هذا اللاتناظر يُفسد التحسين التقليدي، إذ لا يمكن ببساطة افتراض أن النوى وحدها تحدد الأداء. في الواقع، تعيق هذه المكونات بعضها بعضًا.

كيف عالجت Together AI المشكلة

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

• تخطيط جديد لخطوط المعالجة (pipelining) — خطوط معالجة برمجية تُحقق أقصى تداخل بين عمل النوى التنسورية ووحدات SFU والذاكرة دون فترات توقف.
• محاكاة الدالة الأسية عبر كثيرات الحدود — بدلًا من وحدة SFU البطيئة، يحسب المرور الأمامي الدالة الأسية على وحدات FMA السريعة (fused multiply-add).
• TMEM وصيغة 2-CTA MMA — استخدام ذاكرة الموترات الجديدة tensor memory (بسعة 256 KB لكل SM) مع نمط تعمل فيه كتلتا خيوط (thread blocks) على عملية مصفوفية واحدة، مما يقلل حركة البيانات في الذاكرة المشتركة.

إمكانات Blackwell التي أتاحت ذلك

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

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

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

2-CTA MMA — نمط جديد تعمل فيه كتلتا خيوط في وقت واحد على عملية مصفوفية واحدة. هذا يخفض الحمل على الذاكرة المشتركة إلى النصف ويقلل عدد العمليات الذرية (atomic).

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

تبلغ FlashAttention-4 على Blackwell B200 بصيغة BF16 نحو 1605 TFLOPs/s (بنسبة استغلال 71%). وهذا أسرع بـ 1.3 مرة من cuDNN 9.13 وبـ 2.7 مرة من Triton.

وللتوضيح: 1605 بيتافلوب تقارب نصف الأداء الذروي لـ Blackwell، لكنها تُستخلص من نواة انتباه معقدة تحديدًا.

ماذا يعني ذلك

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

*Meta مصنّفة منظمة متطرفة ومحظورة في روسيا الاتحادية.