هل يتفوق سياق اليوم الكامل على عشر دقائق؟ انتباه فلاش وسؤال طول التسلسل

إليك السؤال الذي توجد هذه المقالة للإجابة عليه: إذا تمكن المحول من متابعة يوم تداول كامل بدقة ثانية واحدة بدلاً من نافذة مدتها عشر دقائق، فهل يمكنه التنبؤ بشكل أفضل؟

حتى وقت قريب لم يكن بإمكانك حتى أن تسأل. احتياجات الاهتمام القياسية $\mathcal{O}(N^2)$ الذاكرة، لذا فإن يومًا مكونًا من 23400 خطوة عند 12 رأسًا في float16 يحتاج إلى ما يقرب من 12.9 جيجابايت لمصفوفة النتائج وحدها - أكثر من معلمات النموذج، وأكثر مما ستوفره لك معظم البطاقات. تم إغلاق السؤال بالحساب قبل أن يختبره أي شخص.

يفتحه انتباه الفلاش (داو وآخرون، 2022). ليس عن طريق تقريب الانتباه — بل يحسب نفس النتيجة بالضبط — ولكن عن طريق إعادة هيكلة الحساب ليكون مدركًا للإدخال والإخراج، مما يقلل من حركة المرور بين مستويات ذاكرة وحدة معالجة الرسومات. هذا هو المحتوى المثير للاهتمام حقًا هنا، وقد تم تناول معظم هذه المقالة حول كيفية عمله: التبليط، والتكرار عبر الإنترنت لـ softmax، و $\Theta(N^2 d^2 / M)$ IO ملزمة وإعادة حساب التمرير للخلف.

لكن الآلية هي التمكين، وليس المطالبة. "السياق الأطول هو الأفضل" هو بيان تجريبي حول الأسواق، والموقف الدائم لهذه المدونة - من Temporal Fusion Transformers ، التي وجدت أن محولات الفانيليا في النماذج المالية المفرطة والنماذج المتكررة ذات المراجعة القصيرة تظل قادرة على المنافسة عند الترددات العالية - تقطع الاتجاه الآخر. لذا تنتهي المقالة بالقياس، وليس بالآلية.

لماذا يرتبط الاهتمام بالذاكرة؟

يحسب الاهتمام $\text{softmax}(\mathbf{Q}\mathbf{K}^\top / \sqrt{d_k})\mathbf{V}$ — البدائية نفسها، في سياق التداول، تمت تغطيتها في محولات الاندماج الزمني للتنبؤ بالمحفظة متعددة الأفق. المشكلة برمتها تكمن في سطر واحد من ذلك: مصفوفة الدرجات المتوسطة $\mathbf{S} = \mathbf{Q}\mathbf{K}^\top/\sqrt{d_k}$ يكون $N \times N$، يتم كتابته في الذاكرة، وإعادة قراءته لـ softmax، ثم كتابته مرة أخرى، وقراءته مرة أخرى للماتمول النهائي - ويجب الاحتفاظ به من أجل الانتشار العكسي.

شدة الانتباه الحسابية هي $\approx \min(d, N)$، لذا فإن حوالي 64 FLOP/byte at $d = 64$ - على يسار نقطة التلال A100. إنه يقع على سقف النطاق الترددي المائل، وليس سقف الحوسبة المسطحة: تقضي وحدة معالجة الرسومات وقتًا أطول في التحرك $\mathbf{S}$ حول من ضرب أي شيء. إن إطار عمل خط السقف الذي يستخدمه هذا - نقطة التلال، والسقف المائل مقابل السقف المسطح، ولماذا يقرر نفس المنطق ما إذا كانت وحدة معالجة الرسومات تستحق الشراء على الإطلاق - تم إنشاؤه بأرقام مقاسة في عندما تدفع وحدة معالجة الرسومات.

التسلسل الهرمي للذاكرة الذي تستغله الخوارزمية

تبلغ سعة ذاكرة SRAM تقريبًا 10x عرض النطاق الترددي و100x زمن استجابة أقل، عند جزء من الألف من السعة. كل ما تفعله Flash Attention يتبع هذه التجارة: التخلي عن السعة، وشراء النطاق الترددي وزمن الوصول. نفس خطوة "إعادة هيكلة الخوارزمية بدلاً من شراء الأجهزة"، والتي تم قياسها باستخدام اختبار خلفي لوحدة المعالجة المركزية، هي سلم سرعة الاختبار الخلفي.

خوارزمية انتباه الفلاش

يقوم Flash Attention بمعالجة الانتباه في المربعات ذات الحجم المناسب لذاكرة SRAM، ولا يتم تحقيقها بالكامل أبدًا $N \times N$ المصفوفة في HBM على الإطلاق.

تقسيم $\mathbf{Q}$ داخل $T_r = \lceil N/B_r \rceil$ كتل الصف و $\mathbf{K}, \mathbf{V}$ داخل $T_c = \lceil N/B_c \rceil$ كتل الأعمدة، مع $B_r, B_c$ تم اختياره بحيث يتناسب البلاط بالإضافة إلى مراكماته مع الشريحة. بالنسبة لكل كتلة استعلام، قم بالتكرار على جميع كتل القيمة الرئيسية:

For each query block Q_i:
    Initialize: O_i = 0, m_i = -inf, l_i = 0    # output, running max, running sum
    For each KV block (K_j, V_j):
        1. Load Q_i, K_j, V_j from HBM to SRAM
        2. Compute S_ij = Q_i @ K_j^T / sqrt(d)  # in SRAM
        3. Compute local max: m_ij = rowmax(S_ij)
        4. Compute P_ij = exp(S_ij - m_ij)        # in SRAM
        5. Compute local sum: l_ij = rowsum(P_ij)
        6. Update running statistics:
           m_new = max(m_i, m_ij)
           l_new = l_i * exp(m_i - m_new) + l_ij * exp(m_ij - m_new)
           O_i = O_i * (l_i * exp(m_i - m_new) / l_new)
                + P_ij @ V_j * (exp(m_ij - m_new) / l_new)
           m_i = m_new, l_i = l_new
    Write O_i to HBM

تكرار softmax عبر الإنترنت

الحيلة التي تجعل التبليط ممكنًا هي online softmax. يحتاج softmax الساذج إلى تمريرين فوق الصف: أحدهما للعثور على الحد الأقصى (للاستقرار العددي)، والآخر للتوسيع والتطبيع. يعد التمرير مرتين على صف واحد ترفض تخزينه أمرًا متناقضًا - لذلك يستمر Flash Attention في تشغيل الإحصائيات وإعادة القياس أثناء تقدمه.

بعد الكتل $1, \ldots, j$:

ويتم تصحيح تراكم الإخراج بنفس العامل:

في كل مرة تقوم فيها كتلة جديدة برفع الحد الأقصى للتشغيل، تتم إعادة قياس المخرجات المتراكمة مسبقًا بأثر رجعي $e^{m^{(j-1)} - m^{(j)}}$ - كما لو كان الحد الأقصى الجديد معروفًا منذ البداية. والنتيجة متطابقة جبريًا مع softmax ثنائي التمرير. في الحساب الدقيق هذا ليس تقريبيا. إنها إعادة الارتباط. (في الدقة المتناهية، يكون هناك مسار تقريب مختلف، وهو ما يهم - راجع التحقق من الدقة أدناه.)

تعقيد الإدخال والإخراج

وهذا هو البيان الرسمي للفوز. ينفذ فلاش الانتباه

HBM يصل إلى أين $M$ هو حجم SRAM، ضد $\Theta(Nd + N^2)$ للتنفيذ القياسي. لاحظ أن $M$ يظهر في المقام: كلما كانت لوحة الكتابة على الرقاقة أكبر، قل عدد الرحلات ذهابًا وإيابًا، ولهذا السبب تم ذكر الخوارزمية من حيث التسلسل الهرمي للذاكرة بدلاً من عدد FLOP. لنموذجي $d = 64$ و $M \approx 100$ KB، فإن النسبة تفضل Flash Attention بنسبة وصول أقل بمقدار 5 إلى 10 مرات تقريبًا.

تمرير للخلف: إعادة الحساب بدلاً من التخزين

يحتاج الانتشار العكسي من خلال الانتباه عادةً إلى $\mathbf{P}$ المصفوفة التي رفضت التمريرة الأمامية الاحتفاظ بها. تنبيه الفلاش ** يعيد حساب ** البلاط من $\mathbf{Q}, \mathbf{K}$ أثناء التمرير للخلف، يتم تخزين الإخراج فقط $\mathbf{O}$ وإحصائيات سوفت ماكس $(m, \ell)$ - كلاهما $\mathcal{O}(N)$، لا $\mathcal{O}(N^2)$. إنه يتاجر بكمية متواضعة من العمليات الحسابية الزائدة عن مصطلح الذاكرة الذي كان يمثل المشكلة برمتها. هذه هي نفس الصفقة مثل نقاط التفتيش المتدرجة، والتي يتم تطبيقها على تفاصيل البلاط داخل مشغل واحد.

FA2: التوازي

احتفظ برنامج Flash Attention 2 (Dao, 2023) بالخوارزمية وأصلح الجدولة:

1. ** عدد أقل من عمليات FLOP غير المكتملة. ** قضى FA1 وقتًا حقيقيًا في إعادة القياس، وإيجاد الحد الأقصى، والأس - العمليات التي تعمل على نوى CUDA، وليس نوى الموتر. يؤجل FA2 إعادة القياس إلى نهاية الحلقة الداخلية.
2. التوازي على طول التسلسل. يتوازى FA1 على الدفعة والرؤوس فقط. FA2 يتوازي أيضًا مع كتل الاستعلام. هذا مهم بشكل خاص لحالة التداول، حيث غالبًا ما يكون لديك تسلسل واحد طويل جدًا لكل أصل وحجم دفعة من 1-4 - وهو بالضبط النظام الذي يؤدي فيه توازي الدُفعة والرأس إلى تجويع وحدة معالجة الرسومات.
3. تقسيم عمل الالتواء. يأخذ كل التفاف مجموعة فرعية مختلفة من كتل الاستعلام بدلاً من تقسيم حساب النتيجة والتقليل عبر الالتواءات، مما يؤدي إلى إزالة تقليل الالتواء المتقاطع.

النتيجة المبلغ عنها: ~70% من حالات الذروة النظرية لـ FLOPs على A100 مقابل ~35% لـ FA1.

FA3: ميكانيكا القادوس

Flash Attention 3 (Dao, Shah, 2024) خاص بهندسة H100:

1. ** تخصص الالتواء غير المتزامن. ** يقوم مسرع الذاكرة الموتر (TMA) الخاص بـ Hopper بنقل HBM → SRAM بشكل غير متزامن. يقوم FA3 بتقسيم الإعوجاج إلى منتجين يصدرون أحمال TMA لكتلة KV التالية ومستهلكين يحوسبةون على الكتلة الحالية، بحيث تختفي حركة البيانات خلف العمليات الحسابية.
2. ** ماتمول و سوفت ماكس معشق ** $\mathbf{Q}_i\mathbf{K}_j^\top$ بالنسبة لكتلة واحدة تعمل على نوى موتر بينما تعمل softmax للكتلة السابقة على نوى CUDA - وحدتان مختلفتان من الأجهزة، متزامنتان حقًا بدلاً من تقسيمهما إلى شرائح زمنية.
3. FP8 مع معالجة غير متماسكة. H100 يقوم بـ FP8 بإنتاجية 2x FP16. تم تدمير انتباه FP8 الساذج من خلال القيم المتطرفة. يقوم FA3 بتدوير المتجهات بشكل عشوائي قبل التكميم من حيث الكتلة لنشر الحجم الخارجي عبر الإحداثيات، ويتم الإبلاغ عنه عند خطأ رقمي أقل بمقدار 2.6x من FP8 الساذج.

الإخفاء السببي هو المكان الذي يحصل فيه التداول على خصم

يعد الإخفاء السببي إلزاميًا للسلاسل الزمنية - يجب ألا يهتم النموذج بالمستقبل - وتحت التبليط لا يمثل ذلك تكلفة إضافية بل تكلفة موفرة. أي لوحة تكون مفاتيحها بالكامل في المستقبل بالنسبة لاستعلاماتها يتم تخطيها تمامًا، ولا يتم تحميلها مطلقًا ولا يتم حسابها أبدًا، مما يؤدي إلى قطع نصف العمل تقريبًا. في PyTorch هذا هو is_causal=True; لا شيء آخر مطلوب.

التكامل ثمانية أسطر

لا يتعلق أي من التعليمات البرمجية التي تحتاجها تقريبًا بـ Flash Attention. قم بتبديل مسار مصفوفة النتائج الصريحة للنواة المندمجة:

scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_head)
scores = scores.masked_fill(causal_mask, float("-inf"))
out = torch.matmul(torch.softmax(scores, dim=-1), v)

out = torch.nn.functional.scaled_dot_product_attention(
    q, k, v,
    dropout_p=self.dropout if self.training else 0.0,
    is_causal=True,
)

هذا هو التغيير كله. q, k, v تتشكل (batch, heads, seq, head_dim); لقد اختفى القناع السببي لأن النواة هي التي تبنيه. للحصول على نموذج تداول PyTorch مشروح كامل لإسقاط هذا في - إسقاط الإدخال، والكتل، ورأس من 3 فئات لأعلى/مسطح/لأسفل - استخدم DeepLOB، وللاطلاع على مسار التدريب الكامل، راجع محولات الاندماج الصدغي. إن بناء نسخة رابعة من تلك السقالات هنا لن يعلمنا شيئًا.

المتطلبات: القدرة على الحوسبة >= 8.0 (A100، H100، RTX 3090+)، مدخلات نصف الدقة، PyTorch >= 2.0. تحقق من المسار السريع الذي تم التعامل معه بالفعل torch.backends.cuda.sdp_kernel التشخيص و torch.cuda.max_memory_allocated() - يعود SDPA بصمت إلى نواة الرياضيات في حالة فشل أي شرط مسبق، ويبدو التراجع الصامت تمامًا مثل نموذج العمل الذي يكون بطيئًا فقط.

القياس: هل يفيد السياق الأطول؟

كل ما ورد أعلاه يشير إلى أن سياق 32 كيلو بايت أو 128 كيلو بايت أصبح الآن بأسعار معقولة. لا يذكر شيئًا عما إذا كان مفيدًا أم لا. التجربة الصادقة:

تدريب نفس الهندسة المعمارية في $N \in \{512, 4096, 32768\}$ مع SDPA-Flash على سلسلة BTC المستخدمة في أي مكان آخر في هذه السلسلة، مع الاحتفاظ بالمعلمات والمحسن والهدف ثابتًا بحيث يكون طول التسلسل هو المتغير الوحيد. الإبلاغ عن شيئين:

1. التكلفة. قياس ساعة الحائط لكل عصر و torch.cuda.max_memory_allocated() في كل $N$.
2. الفائدة. الأداء التنبؤي خارج العينة مقابل $N$، في انقسام المشي إلى الأمام.

حملت مسودة سابقة من هذه المقالة جدولًا لأرقام الذاكرة لكل طول تسلسل مشتق تحليليًا من $\mathcal{O}(N)$ صيغة التنشيط. تمت إزالة هذه الصفوف: لم يتم قياسها مطلقًا، كما أنها اختلفت مع حساب ميزانية الذاكرة الخاص بالمقالة. الرقم المشتق المقدم في جدول النتائج هو نتيجة ملفقة، وهذه المدونة لا تقوم بشحنها.

الخاصية المثيرة للاهتمام في هذه التجربة هي أنها قابلة للنشر في أي من الاتجاهين. إذا ارتفع الأداء خارج العينة بشكل رتيب مع $N$، وهذا يبرر برنامج السياق الطويل بأكمله. إذا استقرت على بضعة آلاف من الخطوات، أو انحدرت، فهذه قطعة أقوى - رفيقة لـ السلبية الصادقة - وهذا يعني أن جدار الذاكرة لم يكن أبدًا القيد الملزم على محولات التداول.

المزيد من السياق يعني المزيد من السعة، وبالتالي المزيد من السطح الملائم

هناك سبب محدد لتوقع النتيجة الثابتة أو السلبية. محولات الاندماج الزمني يوثق بالفعل أن محولات الفانيليا طبقت بسذاجة على التراكب في السلسلة المالية - فهي تفتقر إلى التحيزات الاستقرائية الزمنية، وتظل النماذج المتكررة ذات المراجعة القصيرة قادرة على المنافسة عند الترددات العالية. توسيع السياق من 512 إلى 32,768 خطوة لا يضيف معلومات متناسبة مع الطول؛ إن الفارق الهامشي البالغ 32 ألفًا لسلسلة الأسعار شبه الفعالة يحمل القليل جدًا. ما تضيفه بشكل موثوق هو قيمة المعلمات التي تناسبها.

لذا اكتسح $N$ يجب التعامل معه على ما هو عليه: بحث لاختيار النموذج، بنفس الآلية التي تطبقها هذه المدونة على كل بحث آخر. ثلاثة أطوال متتالية مضروبة في أي شيء آخر يختلف هو عدد التجارب، ويجب على الفائز مسح نسبة شارب المفرغة المحسوبة مقابل عدد التجارب هذا وبوابة PBO، وليس مجرد التغلب على جيرانه. بخلاف ذلك، لا يمكن تمييز "انتصارات السياق الطويل" عن اختيار الأفضل من بين ثلاث جولات صاخبة.

التحقق من الدقة، لأن كلمة "دقيق" تعني القيام بالكثير من العمل

تنبيه الفلاش دقيق في الحساب الدقيق. التوصية المرفقة به — تشغيل في FP16 أو BF16، وعلى H100 النظر في FP8 — ليست كذلك. هذه مطالبات منفصلة والثاني هو المهيمن في الممارسة العملية: إعادة ربط المبلغ وخفض الدقة إلى نصفها يعدان من الاضطرابات، ويجب ألا تقوم المقالة التي قدمت ضمان الطلب بعد ذلك بتلويح الدقة يدويًا.

المدونة لديها بالفعل الأداة المناسبة. فخ دقة وحدة معالجة الرسومات يضع المعيار: الدقة المنخفضة لا تحذرك، بل ترجع أشياء غير مقبولة، وتثبت الصحة باستخدام أوراكل التكافؤ على كمية منفصلة في اتجاه مجرى النهر - يتم احتساب التجارة - وليس من خلال منحنيات مقلة العين. تطبق هنا:

• حساب الانتباه باستخدام SDPA-Flash في bf16 ومع تطبيق مرجعي fp64 على مدخلات متطابقة؛ قم بالإبلاغ عن أقصى خطأ نسبي في موتر الإخراج.
• انتقل إلى القرار: بالنسبة للنموذج الذي يُصدر تصنيفًا لأعلى/مسطح/لأسفل، قم بالإبلاغ عن عدد التسميات التي تنقلب بين المسارين، كجزء من إجمالي القرارات.

الخلاف الصغير والمحدود والقابل للتفسير هو توقيع طريق سريع صحيح. تعني التوصية غير المحدودة أن توصية FP8 لم تكن آمنة أبدًا لهذا النموذج. ولا يُعرف أي رقم حتى يتم تشغيله.

متى تصل إليه

مضغوط على القرار، الذي له نفس شكل دليل قرار وحدة معالجة الرسومات:

• أكثر من 2K تقريبًا من الخطوات الزمنية على وحدة معالجة الرسومات CUDA: نعم، دون قيد أو شرط. إنه تغيير من سطر واحد ينتج عنه مخرجات دقيقة، وينمو الفوز مع $N$. لا يوجد سيناريو حيث تريد أن تتحقق $N \times N$ المسار بدلا من ذلك.
• **أقل من 512 خطوة زمنية تقريبًا، على وحدة المعالجة المركزية (CPU)، أو مع البنى التي لا تتطلب الانتباه (شبكات CNN، وأجهزة SSM مثل Mamba): غير ذات صلة. ** على يسار التلال، تمثل النفقات العامة الثابتة التكلفة بأكملها ولم يكن الاهتمام أبدًا عنق الزجاجة.
• الحدود المذكورة أعلاه هي فولكلور، وليست قياسًا — فهي مستمدة من الأدبيات العامة، والتقاطع في النموذج الخاص بك وبطاقتك هو معيار مكون من عشرة أسطر. قم بتشغيله بدلاً من الوثوق بالأرقام المستديرة.

الخلاصة

يعد Flash Attention نتيجة نظيفة ومهمة حقًا: من خلال احترام التسلسل الهرمي للذاكرة وإعادة ربط softmax، فإنه يحسب الانتباه الدقيق باستخدام $\mathcal{O}(N)$ الذاكرة بدلا من $\mathcal{O}(N^2)$، و $\Theta(N^2 d^2 / M)$ IO ملزمة يشرح بالضبط السبب. إن اعتماده في محول التداول هو تغيير سطر واحد بدون تكلفة دقة في الحساب الدقيق وفوز كبير بالذاكرة.

ما لا يفعله هو الإجابة على السؤال الموجود في الأعلى. إنه يحول عبارة "سياق يوم كامل مستحيل" إلى "سياق يوم كامل رخيص"، وهو تغيير في تكلفة التجربة، وليس نتيجتها. إن سقوط جدار الذاكرة هو دعوة للقياس، والقياس هو ما يحول ذلك من ملخص ورقي إلى نتيجة.

المراجع

1. Dao, T., Fu, D.Y., Ermon, S., Rudra, A., Re, C. "FlashAttention: انتباه دقيق سريع وفعال في الذاكرة مع وعي الإدخال والإخراج." نيوريبس (2022). ارشيف:2205.14135
2. داو، ت. "FlashAttention-2: انتباه أسرع مع توازي أفضل وتقسيم العمل." ICLR (2024). arXiv:2307.08691
3. Dao, T., Shah, J. "FlashAttention-3: انتباه سريع ودقيق مع عدم التزامن والدقة المنخفضة." نيوريبس (2024). arXiv:2407.08608
4. فاسواني، أ، وآخرون. "الاهتمام هو كل ما تحتاجه." نيوريبس (2017).
5. Milakov، M.، Gimelshein، N. "حساب التسوية عبر الإنترنت لـ softmax." أرخايف:1805.02867 (2018).