Zyphra ZYFA: FP8 FlashAttention لمعالجات AMD MI355X، أسرع حتى 1.76 مرة

أطلقت Zyphra ZYFA (Zyphra FlashAttention on AMD)، وهي مجموعة من نوى الانتباه (attention kernels) من نوع FP8 وBF16 مخصصة لمرحلة التعبئة المسبقة (prefill) ومعالجات AMD Instinct MI355X. تعمل ZYFA كمحرك الانتباه FP8 الخاص بـ Zyphra للسياقات الطويلة في مرحلة التعبئة المسبقة، والكود مفتوح المصدر بموجب Apache 2.0. في التطبيقات العملية، تُظهر نوى FP8 في ZYFA تحسينات في الأداء تتراوح بين 1.16 و1.76 مرة مقارنةً بـ AITER، مكتبة النوى القياسية من AMD. عند بُعد رأس (head dimension) قدره 128، تكون معدلات الخطأ في FP8 لدى ZYFA أقل بنسبة 62-78% من تلك الخاصة بـ AITER.

باختصار

• الإطار: مكتبة نوى للانتباه الأمامي السببي (causal forward attention) في مرحلة التعبئة المسبقة بدقة FP8 e4m3 وBF16، وبُعد رأس 128، وشكل MLA (192/128).
• يعمل على: حصريًا على AMD Instinct MI355X (CDNA4، gfx950)، وROCm 7.2.3، وPyTorch 2.12. يتميز بواجهة flash_attn_func جاهزة للاستخدام المباشر.
• الأداء: يحقق إنتاجية FP8 تبلغ 1.16-1.76 مرة إنتاجية AITER على التنشيطات الحقيقية؛ أداء BF16 مماثل (1.01-1.02 مرة).
• الأفضل: 2,672 تيرافلوب/ثانية FP8 عند 32 ألف رمز (token)، ليصل إلى 53% من الذروة المعلنة في ورقة بيانات FP8 (AITER: 39%).
• الأسوأ: تحسن بنسبة 1.16 مرة فقط على تنشيطات Mistral-7B؛ سرعة ساعة FP8 تنخفض إلى 1,628 ميغاهرتز مقارنةً بـ 2,108 ميغاهرتز لدى AITER.

ما هي ZYFA؟

تأتي ZYFA مصحوبة بورقة بحثية والكود المرتبط بها من Zyphra. تشغّل محرك الانتباه في Zyphra Optimization Engine (ZOE)، الذي يعمل داخل Zyphra Cloud. نوى BF16 منفذة بلغة HIP، بينما نوى FP8 هي لغة تجميع مجدولة يدويًا مولدة بأداة Python.

لماذا يكون أداء انتباه FP8 ضعيفًا على MI355X؟

تعمل تعليمات المصفوفات FP8 بضعف الأداء الأقصى لـ BF16. لكن FP8 يسرّع فقط ضربي المصفوفات QKᵀ وPV، بينما تعمل عملية softmax بينهما على وحدة المتجهات بالسرعة نفسها لكلتا الدقتين.

من أجل 32 استعلامًا مقابل 128 مفتاحًا، تتطلب معالجة مصفوفات FP8 1,024 دورة، بينما يتطلب softmax بنظام FP8 896 دورة متجهية. وبالتالي، يجب أن تحقق النواة تداخلًا شبه كامل لإبقاء وحدات المصفوفات مستغلة بالكامل. تشير تتبعات الخيوط (thread traces) لدى Zyphra إلى أن خط أنابيب المصفوفات في AITER يكون خاملًا لنسبة 31.7% من الدورات.

تفتقر MI355X إلى تتبع عتادي لاكتمال تعليمات المصفوفات، ما يستلزم أن يخلق ترتيب التعليمات هذا التداخل.

كيف تعمل ZYFA؟

تستخدم ZYFA ثلاث تقنيات:

Lean P: ضرب-جمع واحد بدلًا من دالة أسية

باستخدام نمط بتات e4m3 يُفسَّر كعدد صحيح، تزداد القيمة بمقدار 8 مع كل مضاعفة. الكود الخاص بـ 2^x هو تقريبًا 8x زائد ثابت. تحسب Lean P كل احتمال FP8 باستخدام عملية ضرب-جمع واحدة وتحويل صحيح واحد، ما يقلل دورات المتجهات لعملية softmax لكل كتلة من 896 إلى 384.

جدولة بينغ بونغ بموجتين

تتشارك موجتان كل وحدة SIMD. تصدر إحدى الموجتين تعليمات المصفوفات بينما تحسب الأخرى الاحتمالات، ثم تتبادلان الأدوار. مع Lean P، يبقى خط أنابيب المصفوفات مشغولًا لنسبة 86% من دورات الحالة المستقرة، بزيادة عن 60% مع P الدقيق.

تحجيم لكل كتلة ومجموع صف مقرّب

تُحجَّم كل كتلة من احتمالات 64 مفتاحًا بأس للعدد 2 قبل التقريب. تلغي تعليمة المصفوفات المحجمة في CDNA4 هذا العامل دون تكلفة إضافية. تقسم ZYFA أيضًا على مجموع الاحتمالات المقرّبة، والذي يُحسب بتعليمة مصفوفات إضافية واحدة لكل كتلة.

لمزيد من التفاصيل، توفر Zyphra شرحًا تفاعليًا على موقعها الإلكتروني.

رابط ذو صلة: zyphra.com

المصدر: MarkTechPost