Mixture of Block Attention (MoBA) (Lu et al., 2025) is a promising building block for efficiently processing long contexts in LLMs by enabling queries to sparsely attend to a small subset of key-value blocks, drastically reducing computational cost. However, the design principles governing MoBA's performance are poorly understood, and it lacks an efficient GPU implementation, hindering its practical adoption. In this paper, we first develop a statistical model to analyze MoBA's underlying mechanics. Our model reveals that performance critically depends on the router's ability to accurately distinguish relevant from irrelevant blocks based on query-key affinities. We derive a signal-to-noise ratio that formally connects architectural parameters to this retrieval accuracy. Guided by our analysis, we identify two key pathways for improvement: using smaller block sizes and applying a short convolution on keys to cluster relevant signals, which enhances routing accuracy. While theoretically better, small block sizes are inefficient on GPUs. To bridge this gap, we introduce FlashMoBA, a hardware-aware CUDA kernel that enables efficient MoBA execution even with the small block sizes our theory recommends. We validate our insights by training LLMs from scratch, showing that our improved MoBA models match the performance of dense attention baselines. FlashMoBA achieves up to 14.7x speedup over FlashAttention-2 for small blocks, making our theoretically-grounded improvements practical. Code is available at: https://github.com/mit-han-lab/flash-moba.
تتناول هذه الورقة تحسيناً منهجياً لآلية خليط انتباه الكتل (MoBA). يعالج MoBA السياقات الطويلة بكفاءة من خلال السماح للاستعلامات بالانتباه بشكل متفرق إلى عدد قليل من كتل المفاتيح والقيم، لكن مبادئ تصميمه غير واضحة وتفتقر إلى تنفيذ فعال على وحدات معالجة الرسومات. يؤسس المؤلفون نموذجاً إحصائياً لتحليل آلية MoBA، ويشتقون صيغة نسبة الإشارة إلى الضوضاء SNR ∝ √(d/B)، مما يكشف العلاقة بين معاملات العمارة ودقة الاسترجاع. بناءً على التحليل النظري، يقترحون مسارين للتحسين: استخدام أحجام كتل أصغر وتطبيق التفافات قصيرة على المفاتيح لتجميع الإشارات ذات الصلة. لحل مشكلة عدم الكفاءة على وحدات معالجة الرسومات للكتل الصغيرة، يطورون نواة CUDA تدرك الأجهزة FlashMoBA، مما يحقق تسريعاً يصل إلى 14.7 مرة مقارنة بـ FlashAttention-2، مما يجعل التكوين الأمثل نظرياً قابلاً للتطبيق عملياً.
تتوسع نماذج اللغة الكبيرة (LLMs) إلى مجالات متعددة الأنماط مثل فهم الفيديو والتوليد، مما يتطلب معالجة سياقات فائقة الطول. ومع ذلك، فإن التعقيد الحسابي التربيعي لآلية الانتباه الذاتي يشكل اختناقاً. تحاول طرق الانتباه المتفرق حل هذه المشكلة من خلال الانتباه فقط إلى المناطق المهمة، حيث يعتبر MoBA طريقة واعدة تقلل التعقيد إلى شبه خطي من خلال توجيه كل استعلام إلى عدد قليل من كتل المفاتيح والقيم باستخدام جهاز توجيه قابل للتعلم.
مع توسع نماذج اللغة الكبيرة إلى تطبيقات مثل فهم الفيديو ومعالجة المستندات الطويلة، قد يصل طول السياق إلى ملايين الرموز. يجعل التعقيد O(N²) للانتباه الكثيف هذه التطبيقات غير قابلة للتطبيق حسابياً. آلية انتباه متفرقة فعالة هي تقنية أساسية لتحقيق هذه الرؤية.
على الرغم من أن MoBA جذاب نظرياً، إلا أنه يواجه مشكلتين حاسمتين:
مبادئ التصميم غير واضحة: يفتقر الفهم النظري لكيفية اختيار جهاز التوجيه لعدد قليل من الكتل الصحيحة من آلاف الكتل المرشحة (مشكلة "البحث عن إبرة في كومة قش")
نقص التنفيذ الفعال: خاصة بالنسبة لأحجام الكتل الصغيرة، التنفيذ الأصلي غير فعال، بل أبطأ من الانتباه الكثيف
يعتقد المؤلفون أن هناك حاجة للاختراق على المستويين النظري والعملي: فهم آلية عمل MoBA نظرياً، وتطوير تنفيذ فعال على وحدات معالجة الرسومات لجعل التكوين الأمثل نظرياً قابلاً للتطبيق على الأجهزة.
نموذج نظري إحصائي: تأسيس نموذج إحصائي لآلية اختيار كتل MoBA، واشتقاق صيغة نسبة الإشارة إلى الضوضاء SNR = Δμ_eff√(d/2B)، مما يربط رسمياً معاملات العمارة (d, B) بدقة استرجاع جهاز التوجيه
مبادئ التصميم: بناءً على التحليل النظري، اقتراح والتحقق من مسارين للتحسين:
تحسين نسبة بُعد الرأس إلى حجم الكتل (d/B)، والتحكم في سعة النموذج من خلال تغيير حجم الكتل B
تطبيق التفافات قصيرة على المفاتيح لتحسين تجميع الإشارات
نواة FlashMoBA: تطوير نواة CUDA تدرك الأجهزة، مما يجعل حجم الكتل الصغير الأمثل نظرياً قابلاً للتطبيق عملياً، مما يحقق:
تسريع يصل إلى 14.7 مرة مقارنة بـ FlashAttention-2 لتكوينات الكتل الصغيرة
تسريع 7.4 مرات وتوفير ذاكرة 6.1 مرات مقارنة بتنفيذ MoBA الأصلي عند طول تسلسل 64K
التحقق التجريبي: التحقق من نموذج MoBA المحسّن من خلال التدريب من الصفر لنماذج اللغة، مما يثبت أنه يطابق أداء خط أساس الانتباه الكثيف مع الحفاظ على تفرق 7/8
الإدخال: أزواج المفاتيح والقيم (K, V) واستعلامات Q بطول تسلسل N
الإخراج: إخراج الانتباه O = softmax(QK^T/√d)V
القيد: تقليل التعقيد من O(N²) إلى O(N·kB) من خلال الانتباه المتفرق، حيث k≪n=N/B
يقسم MoBA N مفتاح إلى n=N/B كتل بحجم B. لكل استعلام q، بدلاً من الانتباه إلى جميع N مفاتيح القيم، يتم اختيار فقط أفضل k كتل الأكثر صلة.
1. اختيار Top-K المقسم (Flash TopK)
خط أنابيب ثلاثي المراحل:
المرحلة 1: نواة Triton تحسب مراكز كتل المفاتيح، تولد مصفوفة أصغر K̃
المرحلة 2: نواة مقسمة مستوحاة من FlashAttention-2، تحسب الدرجات بين Q و K̃، تجد أفضل k كتل مفاتيح لكل استعلام، بدون تجسيد مصفوفة الدرجات الكاملة (الخوارزمية 3)
المرحلة 3: إعادة تنسيق فعالة لمؤشرات مراكز الاستعلام إلى تخطيط varlen لمراكز كتل المفاتيح
2. المسار الأمامي: الجمع والتكثيف (الخوارزمية 1)
لكل كتلة استعلام منطقية Q_i:
لكل كتلة مفتاح منطقية K_j:
استخدام مؤشرات varlen للعثور على الاستعلامات ذات الصلة
معالجة دفعة فرعية من الاستعلامات كبلاطات فيزيائية كثيفة:
- جمع بلاطات الاستعلام الفيزيائية من HBM إلى SRAM
- تخزين مؤقت في SRAM، إعادة استخدام عبر جميع بلاطات K_j المنطقية
- تنفيذ ضرب مصفوفات كثيف فعال
- تشتيت النتائج مرة أخرى إلى HBM
التحسين الرئيسي: من خلال تخزين بلاطات الاستعلام المجمعة مؤقتاً في SRAM، إعادة الاستخدام عبر عمليات ضرب مصفوفات كثيفة متعددة، مما يوزع بشكل فعال تكلفة عملية الجمع غير المنتظمة
3. المسار العكسي: إعادة الحساب (الخوارزمية 5)
اعتماد التصميم الفعال للذاكرة من FlashAttention-2
التوازي عبر بُعد المفتاح، كل كتلة خيط تعالج كتلة مفتاح واحدة
عكس استراتيجية "الجمع والتكثيف" من المسار الأمامي
إعادة حساب درجات الانتباه لتجنب تخزين مصفوفة الانتباه الكاملة
استخدام إضافة ذرية إلى مخزن مؤقت عام عالي الدقة لتجميع تدرجات الاستعلام الجزئية (dQ) بأمان
التفاف 1-D سببي قابل للفصل بعمق: groups=hidden_size، تصفية مستقلة لكل قناة
البنية السببية: ملء يساري، الحفاظ على الخاصية الانحدارية الذاتية
حجم النواة: W ∈ {3, 5} (kconv3 و kconv5)
التفعيل والبقايا: تفعيل SiLU + اتصال البقايا
الصيغة:
k'_t = k_t + SiLU(Σ_{ℓ=0}^{W-1} W_ℓ ⊙ k_{t-ℓ})
التأثير: تشجيع تدفق التدرجات بين الرموز المجاورة داخل الكتل أثناء التدريب، مما يعزز ضمنياً محاذاة الرموز المجاورة مع اتجاه الاستعلام، مما يزيد من عدد الرموز ذات الصلة داخل الكتلة m والألفة المتوسطة μ_cluster
عبر عدة معايير وأحجام، MoBA يطابق أو يتفوق على الانتباه الكثيف:
حجم النموذج
المهمة
كثيف
أفضل MoBA
التحسن
340M
دقة LM
44.2%
46.2% (kconv5)
+2.0%
340M
RULER
42.0%
63.9% (kconv5)
+21.9%
340M
LongBench
11.3
13.7 (kconv3)
+2.4
1B
دقة LM
50.9%
52.7% (kconv3)
+1.8%
1B
RULER
61.3%
68.2% (kconv3)
+6.9%
الرؤى الرئيسية:
الانتباه الكثيف يفشل تماماً عند طول 32K (0%)، MoBA-128+kconv5 يصل إلى 100% استرجاع عند 64K
التوجيه المتفرق يخفف من تخفيف الانتباه: مع نمو طول التسلسل، يوزع softmax الكثيف كتلة الاحتمالية على جميع الرموز، بينما MoBA يركز على عدد قليل من الكتل المستهدفة
طرق النمط الثابت: Sparse Transformer (Child et al., 2019), Longformer (Beltagy et al., 2020), BigBird (Zaheer et al., 2021)
الطرق المتعلمة: Reformer (LSH, Kitaev et al., 2020), Linformer (الإسقاط, Wang et al., 2020), Routing Transformer (Roy et al., 2021), Performer (Choromanski et al., 2021)
تحسينات التنفيذ: FlashAttention (Dao et al., 2022; 2023) يحسن IO لكن لا يقلل التعقيد
التحديات: الأنماط المتفرقة والوصول غير المنتظم للذاكرة يصعب تنفيذها بكفاءة
الأدوات: Triton (Tillet et al., 2019) يبسط تطوير النوى، لكن الأداء القصوى تتطلب تحسين دقيق
التحسينات ذات الصلة: FlashDecoding++ (Hong et al., 2024), PagedAttention (Kwon et al., 2023), Ring Attention (Liu et al., 2023), FlashInfer (Ye et al., 2025)
الفرق في هذه الورقة: FlashMoBA محسّنة خصيصاً لنمط الكتل المتفرقة الصغيرة، مما يجعل التكوين الأمثل نظرياً عملياً
ورقة MoBA الأصلية: Lu et al. (2025) - تقديم مفهوم خليط انتباه الكتل
سلسلة FlashAttention: Dao et al. (2022), Dao (2023) - أساس تنفيذ الانتباه الفعال من حيث IO
التفاف المفاتيح: Yang et al. (2025) - قاعدة دلتا للتحويلات الخطية المتوازية
معايير التقييم:
RULER: Hsieh et al. (2024) - تقييم استرجاع السياق الطويل
LongBench: Bai et al. (2024) - فهم السياق الطويل متعدد المهام
طرق متفرقة ذات صلة:
Block Sparse Attention: Guo et al. (2024)
XAttention: Xu et al. (2025)
BigBird: Zaheer et al. (2021)
التقييم الشامل: هذه ورقة ممتازة تجمع بين النظرية والممارسة بشكل وثيق. نظرياً، نموذج SNR يوفر إرشادات واضحة لتصميم الانتباه المتفرق؛ عملياً، FlashMoBA يحول الرؤى النظرية إلى تحسينات أداء حقيقية. على الرغم من القيود في حجم النموذج ونطاق التجارب، فإن المساهمات الأساسية - مبادئ التصميم الرسمية والتنفيذ الفعال - لها أهمية كبيرة لتطور نماذج اللغة الكبيرة بسياق طويل. يستحق الثناء بشكل خاص الموقف الصارم للمؤلفين في التحقق من النظرية من خلال تجارب متغيرات مضبوطة، بالإضافة إلى جهودهم في نشر الكود لتعزيز اعتماد المجتمع.