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 प्रश्नों को कुछ कुंजी-मान ब्लॉकों पर विरल रूप से ध्यान केंद्रित करके लंबे संदर्भ को कुशलतापूर्वक संसाधित करता है, लेकिन इसके डिजाइन सिद्धांत अस्पष्ट हैं और GPU कार्यान्वयन में कमी है। लेखक MoBA तंत्र का विश्लेषण करने के लिए एक सांख्यिकीय मॉडल स्थापित करते हैं, संकेत-से-शोर अनुपात (SNR) सूत्र SNR ∝ √(d/B) प्राप्त करते हैं, जो आर्किटेक्चर पैरामीटर और पुनर्प्राप्ति सटीकता के बीच संबंध को प्रकट करता है। सैद्धांतिक विश्लेषण के आधार पर, दो सुधार पथ प्रस्तावित किए जाते हैं: छोटे ब्लॉक आकार का उपयोग और संबंधित संकेतों को क्लस्टर करने के लिए कुंजियों पर छोटा कनवल्शन लागू करना। GPU पर छोटे ब्लॉकों की दक्षता समस्या को हल करने के लिए, FlashMoBA हार्डवेयर-जागरूक CUDA कर्नल विकसित किए गए हैं, जो FlashAttention-2 की तुलना में 14.7 गुना तक त्वरण प्राप्त करते हैं, जिससे सैद्धांतिक रूप से इष्टतम कॉन्फ़िगरेशन व्यावहारिक रूप से व्यवहार्य हो जाता है।
बड़े भाषा मॉडल (LLMs) वीडियो समझ और पीढ़ी जैसे बहु-मोडल क्षेत्रों में विस्तारित हो रहे हैं, जिन्हें अति-लंबे संदर्भ को संसाधित करने की आवश्यकता है। हालांकि, स्व-ध्यान तंत्र की द्विघात कम्प्यूटेशनल जटिलता एक बाधा बन गई है। विरल ध्यान विधियां केवल महत्वपूर्ण क्षेत्रों पर ध्यान केंद्रित करके इस समस्या को हल करने का प्रयास करती हैं, जिनमें MoBA एक आशाजनक विधि है, जो प्रत्येक प्रश्न को कुछ कुंजी-मान ब्लॉकों की ओर निर्देशित करने के लिए एक राउटर सीखता है, जटिलता को लगभग रैखिक तक कम करता है।
जैसे-जैसे LLMs वीडियो समझ, लंबे दस्तावेज़ प्रसंस्करण आदि अनुप्रयोगों में विस्तारित होते हैं, संदर्भ की लंबाई लाखों टोकन तक पहुंच सकती है। पारंपरिक घने ध्यान की O(N²) जटिलता इन अनुप्रयोगों को कम्प्यूटेशनल रूप से अव्यावहारिक बनाती है। कुशल विरल ध्यान तंत्र इस दृष्टिकोण को साकार करने के लिए मुख्य तकनीक है।
MoBA सैद्धांतिक रूप से आकर्षक है, लेकिन दो महत्वपूर्ण समस्याओं का सामना करता है:
डिजाइन सिद्धांत अस्पष्ट: राउटर हजारों उम्मीदवार ब्लॉकों में से कुछ सही ब्लॉकों को कैसे विश्वसनीय रूप से चुनता है ("समुद्र में सुई" समस्या) इसकी सैद्धांतिक समझ की कमी है
कुशल कार्यान्वयन की कमी: विशेष रूप से छोटे ब्लॉक आकार के लिए, मूल कार्यान्वयन अक्षम है, यहां तक कि घने ध्यान से भी धीमा है
लेखक मानते हैं कि सैद्धांतिक और व्यावहारिक दोनों स्तरों पर सफलता की आवश्यकता है: सैद्धांतिक रूप से MoBA के कार्य तंत्र को समझना, व्यावहारिक रूप से कुशल GPU कार्यान्वयन विकसित करना, जिससे सैद्धांतिक रूप से इष्टतम कॉन्फ़िगरेशन हार्डवेयर पर व्यवहार्य हो।
सांख्यिकीय सिद्धांत मॉडल: MoBA ब्लॉक चयन तंत्र का एक सांख्यिकीय मॉडल स्थापित किया, संकेत-से-शोर अनुपात सूत्र SNR = Δμ_eff√(d/2B) प्राप्त किया, जो आर्किटेक्चर पैरामीटर (d, B) को राउटर पुनर्प्राप्ति सटीकता से औपचारिक रूप से जोड़ता है
डिजाइन सिद्धांत: सैद्धांतिक विश्लेषण के आधार पर दो सुधार पथ प्रस्तावित और सत्यापित किए:
सिर आयाम और ब्लॉक आकार अनुपात (d/B) को अनुकूलित करना, ब्लॉक आकार B को बदलकर मॉडल क्षमता को नियंत्रित करना
संकेत क्लस्टरिंग में सुधार के लिए कुंजियों पर छोटा कनवल्शन लागू करना
FlashMoBA कर्नल: हार्डवेयर-जागरूक CUDA कर्नल विकसित किए, जो सैद्धांतिक रूप से इष्टतम छोटे ब्लॉक आकार को व्यावहारिक रूप से व्यवहार्य बनाते हैं:
छोटे ब्लॉक कॉन्फ़िगरेशन के लिए FlashAttention-2 की तुलना में 14.7 गुना तक त्वरण
64K अनुक्रम लंबाई पर मूल MoBA कार्यान्वयन की तुलना में 7.4 गुना त्वरण और 6.1 गुना मेमोरी बचत
अनुभवजन्य सत्यापन: शुरुआत से LLMs को प्रशिक्षित करके सुधारे गए MoBA मॉडल को सत्यापित किया, जो 7/8 विरलता बनाए रखते हुए घने ध्यान आधारभूत के प्रदर्शन से मेल खाता है
इनपुट: लंबाई N के साथ कुंजी-मान जोड़े (K, V) और प्रश्न Q
आउटपुट: ध्यान आउटपुट 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: Triton कर्नल कुंजी ब्लॉक केंद्रक की गणना करता है, छोटा मैट्रिक्स K̃ उत्पन्न करता है
चरण 2: FlashAttention-2 से प्रेरित tiled कर्नल, Q और K̃ के बीच स्कोर की गणना करता है, प्रत्येक प्रश्न के शीर्ष-k कुंजी ब्लॉक खोजता है, पूर्ण स्कोर मैट्रिक्स को भौतिक रूप दिए बिना (एल्गोरिथ्म 3)
चरण 3: प्रश्न केंद्र सूचकांकों को कुंजी ब्लॉक केंद्रों के varlen लेआउट में पुनः प्रारूपित करने के लिए उच्च दक्षता वाली epilogue
2. फॉरवर्ड पास: Gather-and-Densify (एल्गोरिथ्म 1)
प्रत्येक तार्किक प्रश्न ब्लॉक Q_i के लिए:
प्रत्येक तार्किक कुंजी ब्लॉक K_j के लिए:
varlen सूचकांक का उपयोग करके प्रासंगिक प्रश्न खोजें
प्रश्न उप-समूह को घने भौतिक ब्लॉकों में बैच करें:
- HBM से भौतिक प्रश्न ब्लॉक को SRAM में एकत्रित करें
- SRAM में कैश करें, तार्किक कुंजी ब्लॉक K_j के सभी भौतिक टाइलों में पुनः उपयोग करें
- उच्च दक्षता वाली घनी GEMM निष्पादित करें
- परिणाम को HBM में बिखेरें
मुख्य अनुकूलन: SRAM में एकत्रित प्रश्न ब्लॉकों को कैश करके, कई घनी GEMMs में पुनः उपयोग करके, अनियमित gather ऑपरेशन की लागत को प्रभावी रूप से परिशोधित किया जाता है
3. बैकवर्ड पास: पुनर्गणना (एल्गोरिथ्म 5)
FlashAttention-2 के मेमोरी-कुशल डिजाइन को अपनाता है
कुंजी आयाम में समानांतरकरण, प्रत्येक थ्रेड ब्लॉक एक कुंजी ब्लॉक को संभालता है
फॉरवर्ड प्रोपेगेशन की "gather-and-densify" रणनीति को दर्पण करता है
पूर्ण ध्यान मैट्रिक्स को संग्रहीत करने से बचने के लिए ध्यान स्कोर की पुनर्गणना करता है
आंशिक प्रश्न ग्रेडिएंट (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 Acc
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 Acc
50.9%
52.7% (kconv3)
+1.8%
1B
RULER
61.3%
68.2% (kconv3)
+6.9%
मुख्य अंतर्दृष्टि:
घना ध्यान 32K लंबाई पर पूरी तरह विफल (0%), MoBA-128+kconv5 64K पर 100% तक पहुंचता है
विरल राउटिंग ध्यान को कम करता है: अनुक्रम लंबाई बढ़ने के साथ, घना 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) - Mixture of Block Attention अवधारणा प्रस्तावित करता है
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 सैद्धांतिक अंतर्दृष्टि को वास्तविक प्रदर्शन सुधार में परिवर्तित करता है। मॉडल स्केल और प्रयोगात्मक सीमा के बावजूद, इसके मुख्य योगदान—औपचारिक डिजाइन सिद्धांत और कुशल कार्यान्वयन—लंबे संदर्भ LLMs के विकास के लिए महत्वपूर्ण हैं। विशेष रूप से सराहनीय है सैद्धांतिक सत्यापन के लिए नियंत्रण चर प्रयोग के माध्यम से लेखकों का कठोर दृष्टिकोण, और समुदाय को अपनाने को बढ़ावा देने के लिए खुला स्रोत कोड।