أصدر فريق MoonMath AI نواة الاهتمام الأمامي bf16 لوحدة معالجة الرسوميات MI300X من AMD. إنه مكتوب بلغة HIP، وليس تجميعًا مكتوبًا بخط اليد. الكود مفتوح المصدر بموجب ترخيص MIT. أفاد فريق MoonMath.ai أنه يتفوق على AITER v3، نواة AMD المحسنة، في كل شكل تم اختباره. يأتي الوصول إلى المعدن من HotAisle، وهو مزود سحابة AMD.
الاهتمام هو تنصهر softmax(QKᵀ/√d)·V التشغيل داخل كل محول. MI300X هو وحدة معالجة الرسومات لمركز البيانات CDNA3 من AMD، مع هدف ISA (gfx942). تعمل هذه النواة على هذا الجهاز فقط.
ليرة تركية؛ د
- MoonMath.ai مفتوح المصدر لنواة bf16 للأمام لـ AMD MI300X، مكتوبة بلغة HIP، وليس التجميع (MIT).
- إنه يتفوق على AMD’s AITER v3 في كل شكل ووضع تقريب – Geomean 1.18×/1.15×/1.08×، حتى 1.26×.
- الحيلة الأساسية: تتيح لك أغلفة asm ذات التعليمات الواحدة اختيار كود التشغيل بينما يقوم المترجم بتخصيص السجلات.
- معظم عمليات التسريع هي وضع الذاكرة – K في LDS، وV hot في L1، وQ والمراكم في السجلات.
- استخدمته شركة SGLang PR الحقيقية لتسريع نشر فيديو Wan2.1 بمقدار 1.23×، دون أي تراجع في الجودة.
فهم النواة
النواة عبارة عن برنامج صغير يعمل مباشرة على العديد من مراكز وحدة معالجة الرسومات لإجراء عملية حسابية محددة – هنا، حسابات الانتباه – بالسرعة التي تسمح بها الأجهزة. تحسب النواة الانتباه الأمامي في bf16 على MI300X فقط. فهو يأخذ مدخلات إما في تخطيط BSHD أو BHSD، بدون تبديل. تم تثبيت بُعد الرأس عند 128. وهو يدعم أي طول تسلسلي، بما في ذلك الانتباه المتبادل.
هناك حدود حقيقية. لا يوجد قناع سببي، ولا GQA، ولا يوجد خلط للفارلين. المخرجات هي bf16، وتعمل على أجهزة gfx942 حصريًا.
يتم التحكم في الأرقام بإحكام. تتوافق أوضاع التقريب الثلاثة جميعها مع قاعدة التقريب لكل وضع الخاصة بـ AITER. يقع كل مخرج محدود ضمن 1 bf16 ULP من AITER. التعامل مع NaN وInf متطابقان بعض الشيء، والنتائج حتمية.
الخدعة الأساسية: أغلفة asm ذات التعليمات الواحدة
تتجنب التقنية الأساسية معضلة مألوفة. تحافظ جوهرية المترجم على الكود مرتبًا ولكنها تسمح للمترجم بإعادة ترتيب المعاملات أو إعادة تسميتها. يمنح التجميع المضمن الخام التحكم ولكنه يفرض التسجيل اليدوي وإدارة العناوين.
يقوم MoonMath بتغليف تعليمات واحدة بالضبط في ملف __device__ __forceinline__ وظيفة. تصف قيود asm الموسعة المعاملات. يختار فريق البحث رمز التشغيل. لا يزال المترجم يخصص السجلات ويتتبع تدفق البيانات.
// in/out tied to the SAME VGPR → no accumulator rename, no v_mov copy.
__device__ __forceinline__ void asm_mfma(bf16x4_t a, bf16x4_t b, fp32x4_t& c) {
asm volatile("v_mfma_f32_16x16x16_bf16 %0, %1, %2, %0"
: "+v"(c) : "v"(a), "v"(b));
}ال "+v"(c) يربط القيد مدخلات ومخرجات المجمع بنفس VGPR. لا يتم إصدار أي تعليمات نسخ. هذا يبقي النواة قريبة من HIP العادي. لا يزال يوجه الآلة تعليمة واحدة في كل مرة.
الهندسة المعمارية: ثماني موجات، مجموعتان، حاجزان
تحتوي وحدة حساب CDNA3 على أربع وحدات SIMD. كتلة الكتاب المدرسي هي أربع موجات. بدلاً من ذلك، يقوم MoonMath بتشغيل ثماني موجات لكل كتلة، في مجموعتين من أربعة.
المجموعتان تعملان بنفس الطريقة Q*Kسوفت ماكس, O += P*V تسلسل. يتم تعويضهم بمرحلة. بينما تقوم إحدى المجموعتين بتشبع نواة المصفوفة، تقوم الأخرى بتشغيل softmax وإصدار الأحمال. ثم يتم التبادل، بحيث لا يتوقف قلب المصفوفة أبدًا.
هناك نوعان s_barrierلكل تكرار. واحد يجلس في مرحلة التسليم. واحد يجلس على حدود التكرار. تتعامل عمليات الانتظار لكل عداد مع بقية المزامنة.
هذا يعكس تناوب Matmul و softmax الخاص بـ FlashAttention-3. إنه لا ينسخ تقسيم المنتج والمستهلك الخاص بـ FA3. في CDNA3، تكون كل حركة للذاكرة غير متزامنة بالفعل، لذا فإن موجة المنتج المخصصة غير ضرورية.
أين تعيش البيانات ولماذا 16×16×16
معظم التسريع يأتي من وضع الذاكرة. K التدفقات من HBM إلى LDS، مزدوجة المخزنة، ومشتركة بين جميع الموجات الثمانية. V يبقى ساخنًا في L1، اقرأ على كل matmul PV. Q والمراكم تعيش في السجلات.
اختار فريق البحث 16×16×16 MFMA على 32×32×8. كلا الشكلين لهما إنتاجية متطابقة. يتراكم البلاط الأصغر في 4 عناصر fp32 لكل حارة، مقابل 16. ويترك الضغط المنخفض للمراكم مساحة للجلب المسبق الأعمق والثالث Q بلاط.
| قرار | خيار | سبب |
|---|---|---|
| موجات لكل كتلة | 8 (مجموعتان من 4) | تخطيط خط الأنابيب مباشرة؛ مشاركة نسخة K واحدة |
| شكل مفما | 16×16×16 بف16 | نفس الإنتاجية، وانخفاض ضغط VGPR، وتحسين كفاءة الطاقة |
| التنسيب K | LDS، مخزن مزدوج، 32 كيلو بايت | مشتركة بين جميع الموجات الثمانية، ويتم تبديلها لكل تكرار |
| موضع V | L1، مقيم، محضر مسبقًا | إعادة قراءة عبر الكهروضوئية، وتبقى ساخنة عمدا |
| س + مراكم | VGPRs | قراءة كل التكرار، أبدا إعادة تحميل |
فوزان لاحقان يسدان الفجوة. ثلث Q يؤدي البلاط (3Q) إلى زيادة إعادة استخدام البيانات لكل تحميل K و V بلاط. ينقذ ذيل KV ذو نمط فك التشفير الفلاش الجولة الكسرية العالقة عبر 304 وحدة MI300X. هذه الانتصارات تتالي. تتحرك V لL1 حرر LDS أن الثالث Q البلاط ثم يملأ.
المعيار
تم إجراء الاختبارات على MI300X في bf16، بعد الرأس 128. وتم قياس كل شكل في ثلاثة أوضاع تقريب. يقوم RTNE بتقريب الرقم إلى أقرب رقم متساوي. يقوم RTNA بتقريبه إلى أقرب العلاقات بعيدًا عن الصفر. يقتطع RTZ نحو الصفر.
| الشكل (ب، ح، ق، د) | دائري | لنا (مللي ثانية) | AITER v3 (مللي ثانية) | ضد ايتر | مقابل ماكس |
|---|---|---|---|---|---|
| (2، 24، 8192، 128) | رتني | 3.083 | 3.792 | 1.23× | 1.37× |
| (2، 24، 16384، 128) | رتني | 11.670 | 14.691 | 1.26× | 1.54× |
| (4، 16، 16384، 128) | ار تي زد | 15.055 | 16.183 | 1.07× | 1.47× |
| (2، 24، 32768، 128) | رتنا | 44.440 | 52.363 | 1.18× | 1.57× |
| (1، 16، 131072، 128) | رتني | 232.517 | 269.278 | 1.16× | 1.46× |
Geomeans عبر الاجتياح يفضلون MoonMath. مقابل AITER، يسجل 1.18× (RTNE)، 1.15× (RTNA)، و1.08× (RTZ). مقابل Modular MAX، تتراوح القيم الجغرافية من 1.44× إلى 1.49×، وتصل التسريعات لكل شكل إلى 1.59×.
RTZ هو الوضع الأسرع الخاص بـ AITER وأشد سباق. انتقل شكل RTZ (4، 16، 16384) من 0.95× إلى 1.07×. إن انقسام الذيل KV هو ما أغلق تلك الفجوة النهائية.
الشرح التفاعلي
حالات الاستخدام
يتم تثبيت النواة باستخدام النقطة وتكشف عن واجهة برمجة تطبيقات صغيرة. يتم تشغيله على دفق المتصل، بحيث يتداخل داخل خطوط أنابيب أكبر.
import torch
import moonmath_attention as ma
# PyTorch's ROCm build uses the "cuda" device string on AMD GPUs
q = torch.randn(2, 8192, 24, 128, dtype=torch.bfloat16, device="cuda")
k = torch.randn(2, 8192, 24, 128, dtype=torch.bfloat16, device="cuda")
v = torch.randn(2, 8192, 24, 128, dtype=torch.bfloat16, device="cuda")
out = ma.forward(q, k, v, layout="bshd")
out_rtz = ma.forward(q, k, v, layout="bshd", round_mode="rtz")إحدى حالات الاستخدام الملموسة هي نشر الفيديو. أضاف الفريق دعم LiteAttention وأرسل رسالة عامة إلى SGLang diffusion. في Wan2.1-T2V-1.3B-Diffusers، قاموا بتحويل الاهتمام من AITER إلى liteattention_rocm. تم تحسين الجيل الشامل بمقدار 1.23× على MI300X، مع عدم وجود تراجع واضح في الجودة.
يناسب تخطيط BSHD موترات الانتشار مباشرة. يعمل الانتباه المتقاطع مع أي طول KV وبدون حشوة.
الوجبات السريعة الرئيسية
- النواة هي bf16 إلى الأمام لـ MI300X، مكتوبة في HIP تحت معهد ماساتشوستس للتكنولوجيا.
- إنه يتفوق على AITER v3 في كل شكل ووضع تقريب، بمتوسط جغرافي 1.18×/1.15×/1.08×.
- توفر أغلفة asm ذات التعليمات الواحدة التحكم في كود التشغيل بينما يقوم المترجم بتخصيص السجلات.
- أدى وضع الذاكرة إلى تحقيق معظم المكاسب: K في LDS، وV hot في L1، وQ في السجلات.
- قامت شركة SGLang PR الحقيقية بتسريع نشر فيديو Wan2.1 بمقدار 1.23× دون أي تراجع في الجودة.
تحقق من التفاصيل الفنية. أيضا، لا تتردد في متابعتنا على تغريد ولا تنسى الانضمام إلينا 150 ألف + مل والاشتراك في النشرة الإخبارية لدينا. انتظر! هل أنت على برقية؟ الآن يمكنك الانضمام إلينا على التليجرام أيضًا.
هل تحتاج إلى الشراكة معنا للترويج لصفحة GitHub Repo أو صفحة الوجه المعانقة أو إصدار المنتج أو الندوة عبر الويب وما إلى ذلك؟ تواصل معنا
اكتشاف المزيد من كحيل | أخبار التقنية
اشترك للحصول على أحدث التدوينات المرسلة إلى بريدك الإلكتروني.
