Flash Attention

Flash Attention (FlashAttention) خوارزمية وعائلة من نوى GPU تحسب الانتباه بدقة لنماذج المحوّلات من دون كتابة مصفوفة الانتباه كاملةً إلى الذاكرة الرئيسية لوحدة GPU. ينشئ الانتباه الذاتي القياسي على N من الرموز مصفوفة درجات بحجم N × N، ويكتبها في الذاكرة عالية النطاق الترددي (HBM)، ثم يقرأها مجددًا لحساب softmax وحاصل الضرب في القيم؛ لذا يذهب معظم وقت تنفيذ الانتباه إلى نقل البيانات. قدّم Tri Dao وزملاؤه في Stanford الأسلوب في ورقة عام 2022 بعنوان FlashAttention: انتباه دقيق سريع وموفر للذاكرة مع مراعاة عمليات الإدخال والإخراج. يقسم الأسلوب العمل إلى كتل تتسع لها ذاكرة SRAM السريعة داخل الشريحة، لذا ينمو استخدام الذاكرة خطيًا مع طول التسلسل، وتطابق المخرجات الانتباه القياسي مع مراعاة التقريب بالفاصلة العائمة. تعمل نواه داخليًا في PyTorch وفي محركات استدلال النماذج اللغوية الكبيرة (LLM) مثل vLLM.

آلية عمل Flash Attention#

تحتوي وحدة GPU على مقدار صغير من الذاكرة فائقة السرعة بجوار وحدات الحساب، وذاكرة أكبر وأبطأ خارج الشريحة. تصف ورقة FlashAttention وحدة NVIDIA A100 بذاكرة HBM سعتها 40-80 GB ونطاق ترددي 1.5-2.0 TB/s، وبذاكرة SRAM سعتها 192 KB في كل واحدة من وحدات المعالجة متعددة التدفقات البالغ عددها 108، وبنطاق ترددي مقدّر بـ19 TB/s. تنفذ عملية الانتباه عددًا قليلًا من العمليات لكل بايت تقرأه، لذا تعتمد سرعتها على عدد مرات انتقال البيانات بين المستويين؛ وتحسب الطريقة عمليات النقل هذه، وهذا هو المقصود بعبارة «واعٍ بالإدخال والإخراج».

تحمّل FlashAttention كتلة من الاستعلامات، ثم تمرر كتل المفاتيح والقيم عبر ذاكرة SRAM. وتحسب لكل كتلة الدرجات، وتحدّث الحد الأقصى والمجموع الجاريين لكل صف استعلام، وتعيد تحجيم المخرج الجزئي كلما زاد الحد الأقصى، ثم تضيف مساهمة الكتلة الجديدة. يتيح softmax التزايدي هذا الحصول على نتيجة softmax الدقيقة دون الاطلاع على صف كامل دفعة واحدة، ولا يعود إلى HBM سوى المخرج النهائي ورقمان لكل صف. تعيد المرحلة الخلفية حساب كتل الدرجات من Q وK وV بدلًا من تخزين المصفوفة N × N؛ وهذا يزيد العمليات، لكنه يظل أسرع لأنه يقرأ قدرًا أقل بكثير من الذاكرة. إليك الحلقة نفسها في PyTorch، إذ تعالج جميع الاستعلامات مقابل كتلة واحدة من المفاتيح والقيم في كل مرة:

import torch

torch.manual_seed(0)
n, d, block = 4096, 64, 256
q, k, v = torch.randn(3, n, d).unbind(0)

# Standard attention builds the full n x n score matrix in memory
reference = torch.softmax(q @ k.T / d**0.5, dim=-1) @ v

# FlashAttention-style: stream key/value tiles, keeping a running max and sum per query (online softmax)
row_max, row_sum, out = torch.full((n, 1), -torch.inf), torch.zeros(n, 1), torch.zeros(n, d)
for start in range(0, n, block):
    scores = q @ k[start : start + block].T / d**0.5  # only an n x block tile exists at once
    new_max = torch.maximum(row_max, scores.max(dim=-1, keepdim=True).values)
    rescale, probs = torch.exp(row_max - new_max), torch.exp(scores - new_max)
    row_sum = row_sum * rescale + probs.sum(dim=-1, keepdim=True)
    out = out * rescale + probs @ v[start : start + block]
    row_max = new_max
out = out / row_sum

print(f"Largest score tile: {tuple(scores.shape)} vs full matrix: ({n}, {n})")
print(f"Max difference from standard attention: {(out - reference).abs().max().item():.1e}")

أكبر كتلة درجات حجمها 4,096 × 256، أي سدس عشر مصفوفة 4,096 × 4,096، ولا تختلف النتيجة عن الانتباه القياسي بأكثر من 4.3e-07، وهو مقدار خطأ التقريب في FP32. كما تقسم النواة الفعلية الاستعلامات إلى كتل، واحدة لكل كتلة خيوط في GPU، وتدمج الحلقة كاملةً في نواة واحدة بحيث لا تغادر الكتل ذاكرة SRAM.

FlashAttention-2 و-3 و-4#

يستهدف كل إصدار جيلًا أحدث من وحدات GPU. أعاد FlashAttention-2 (2023) تنظيم توزيع العمل على الخيوط، وبلغت سرعته نحو ضعفي الإصدار الأصلي، محققًا 50-73% من ذروة الإنتاجية في A100. يستخدم FlashAttention-3 (2024) نوى Tensor غير المتزامنة في H100 ويدعم FP8، محققًا 740 TFLOPs/s في FP16 ونحو 1.2 PFLOPs/s في FP8. يستهدف FlashAttention-4 (2026) وحدات GPU من Blackwell، ويحقق وفق التقارير ما يصل إلى 1,613 TFLOPs/s في BF16 على B200، أي 71% من الذروة. تُوزّع النوى ضمن مكتبة flash-attention مفتوحة المصدر، التي تعمل على مدخلات FP16 وBF16، ويختار PyTorch تلقائيًا نواة FlashAttention-2 عبر scaled_dot_product_attention عندما تدعمها وحدة GPU والمدخلات، وإلا فيلجأ إلى تطبيقات موفرة للذاكرة أو عادية. يحسب SAM 2 من Meta، الذي يعمل ضمن حزمة Ultralytics، الانتباه في مُرمّز الصور Hiera وفي انتباه الذاكرة عبر تلك الدالة؛ ولذلك يوجّه PyTorch هذه الطبقات إلى FlashAttention كلما سمحت الأجهزة ودقة المدخلات بذلك.

Flash Attention والمصطلحات ذات الصلة#

• الانتباه القياسي: العملية الرياضية نفسها، مع حساب مصفوفتَي الدرجات والاحتمالات كاملتين وكتابتهما في HBM.
• الانتباه المتناثر: يحسب الدرجات لأزواج مختارة من الاستعلامات والمفاتيح فقط. وقدّمت ورقة FlashAttention أيضًا نسخة متناثرة على مستوى الكتل تتخطى الكتل الفارغة.
• الانتباه الخطي: يستبدل softmax بخريطة ميزات ليزداد التعقيد خطيًا مع الطول، ما يغيّر مخرجات النموذج.
• PagedAttention: مدير ذاكرة يخزّن ذاكرة KV المؤقتة في صفحات ثابتة الحجم أثناء تقديم الخدمة. وتقرأ نوى الانتباه، بما فيها FlashAttention، المفاتيح والقيم من تلك الصفحات.
• الانتباه الحلقي: يقسم تسلسلًا طويلًا جدًا بين الأجهزة، ويحسب كل جهاز الانتباه على مستوى الكتل في الجزء الخاص به، بينما تنتقل كتل المفاتيح والقيم حول حلقة.

الأثر في التدريب والسياقات الطويلة#

بفضل نمو استخدام الذاكرة خطيًا مع طول التسلسل، جعل FlashAttention تدريب نوافذ السياق الطويلة عمليًا. أفادت الورقة الأصلية بتسريع شامل بنسبة 15% لنموذج BERT-large عند طول تسلسل 512 مقارنةً بسجل التدريب MLPerf 1.1، وتسريع بمقدار 3 أضعاف لـGPT-2 عند طول تسلسل 1,024، وبمقدار 2.4 ضعف في معيار Long Range Arena. وتنطبق الوفورات نفسها على الرؤية، حيث يحوّل Transformer للرؤية (ViT) أو Transformer الانتشاري الصور ومقاطع الفيديو عالية الدقة إلى تسلسلات رموز طويلة. كما حسّنت التسلسلات الأطول الدقة: فقد أبلغت الورقة عن أول نماذج Transformer تتجاوز التخمين العشوائي في Path-X، وهي مهمة تختبر ما إذا كان مساران بين نقطتين في صورة 128 × 128، حُوّلت إلى تسلسل من 16,384 بكسل، متصلين.