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 ∝ √(d/B) প্রকাশ করেন, যা স্থাপত্য পরামিতি এবং পুনরুদ্ধার নির্ভুলতার মধ্যে সম্পর্ক প্রকাশ করে। তাত্ত্বিক বিশ্লেষণের উপর ভিত্তি করে, তারা দুটি উন্নতির পথ প্রস্তাব করেন: ছোট ব্লক আকার ব্যবহার করা এবং সম্পর্কিত সংকেত ক্লাস্টার করার জন্য কীতে স্বল্প কনভোলিউশন প্রয়োগ করা। GPU-তে ছোট ব্লকের দক্ষতা সমস্যা সমাধানের জন্য, তারা FlashMoBA হার্ডওয়্যার-সচেতন CUDA কার্নেল বিকাশ করেন, যা FlashAttention-2 এর তুলনায় সর্বোচ্চ ১৪.৭ গুণ ত্বরণ অর্জন করে, তাত্ত্বিকভাবে সর্বোত্তম কনফিগারেশনকে ব্যবহারিকভাবে সম্ভব করে তোলে।
বড় ভাষা মডেল (LLMs) ভিডিও বোঝাপড়া এবং উৎপাদনের মতো মাল্টিমোডাল ক্ষেত্রে সম্প্রসারিত হচ্ছে, যার জন্য অতি-দীর্ঘ প্রসঙ্গ প্রক্রিয়া করার প্রয়োজন। তবে, স্ব-মনোযোগ মেকানিজমের দ্বিঘাত গণনা জটিলতা একটি বাধা হয়ে উঠেছে। স্পার্স অ্যাটেনশন পদ্ধতি শুধুমাত্র গুরুত্বপূর্ণ অঞ্চলে মনোযোগ দিয়ে এই সমস্যা সমাধানের চেষ্টা করে, যেখানে MoBA একটি প্রতিশ্রুতিশীল পদ্ধতি, যা প্রতিটি প্রশ্নকে কম সংখ্যক কী-ভ্যালু ব্লকে পরিচালিত করতে শিখে, জটিলতা প্রায় রৈখিক করে।
LLMs ভিডিও বোঝাপড়া, দীর্ঘ ডকুমেন্ট প্রক্রিয়াকরণ ইত্যাদি অ্যাপ্লিকেশনে সম্প্রসারিত হওয়ার সাথে সাথে, প্রসঙ্গ দৈর্ঘ্য মিলিয়ন-স্তরের টোকেনে পৌঁছাতে পারে। ঐতিহ্যবাহী ঘন মনোযোগের O(N²) জটিলতা এই অ্যাপ্লিকেশনগুলিকে গণনাগতভাবে অসম্ভব করে তোলে। দক্ষ স্পার্স অ্যাটেনশন মেকানিজম এই দৃষ্টিভঙ্গি বাস্তবায়নের জন্য মূল প্রযুক্তি।
MoBA তাত্ত্বিকভাবে আকর্ষণীয় হলেও, দুটি মূল সমস্যার সম্মুখীন:
ডিজাইন নীতি অস্পষ্ট: রাউটার হাজার হাজার প্রার্থী ব্লক থেকে নির্ভরযোগ্যভাবে কম সংখ্যক সঠিক ব্লক নির্বাচন করে কীভাবে ("সমুদ্রে সুই খোঁজা" সমস্যা) তার তাত্ত্বিক বোঝাপড়া অভাব রয়েছে
দক্ষ বাস্তবায়ন অভাব: বিশেষত ছোট ব্লক আকারের জন্য, মূল বাস্তবায়ন অদক্ষ, এমনকি ঘন মনোযোগের চেয়ে ধীর
লেখকরা বিশ্বাস করেন যে তাত্ত্বিক এবং ব্যবহারিক উভয় স্তরে অগ্রগতির প্রয়োজন: তাত্ত্বিকভাবে MoBA-র কাজের মেকানিজম বোঝা, ব্যবহারিকভাবে দক্ষ GPU বাস্তবায়ন বিকাশ করা, তাত্ত্বিকভাবে সর্বোত্তম কনফিগারেশনকে হার্ডওয়্যারে সম্ভব করা।
পরিসংখ্যানগত তত্ত্ব মডেল: MoBA ব্লক নির্বাচন মেকানিজমের একটি পরিসংখ্যানগত মডেল প্রতিষ্ঠা করা, সিগন্যাল-টু-নয়েজ অনুপাত সূত্র SNR = Δμ_eff√(d/2B) প্রকাশ করা, স্থাপত্য পরামিতি (d, B) কে রাউটার পুনরুদ্ধার নির্ভুলতার সাথে আনুষ্ঠানিকভাবে সংযুক্ত করা
ডিজাইন নীতি: তাত্ত্বিক বিশ্লেষণের উপর ভিত্তি করে দুটি উন্নতির পথ প্রস্তাব এবং যাচাই করা:
মাথার মাত্রা এবং ব্লক আকারের অনুপাত (d/B) অপ্টিমাইজ করা, ব্লক আকার B পরিবর্তন করে মডেল ক্ষমতা নিয়ন্ত্রণ করা
সংকেত ক্লাস্টারিং উন্নত করতে কীতে স্বল্প কনভোলিউশন প্রয়োগ করা
FlashMoBA কার্নেল: হার্ডওয়্যার-সচেতন CUDA কার্নেল বিকাশ করা, তাত্ত্বিকভাবে সর্বোত্তম ছোট ব্লক আকারকে ব্যবহারিকভাবে সম্ভব করা, অর্জন করা:
ছোট ব্লক কনফিগারেশনের জন্য FlashAttention-2 এর তুলনায় সর্বোচ্চ ১৪.৭ গুণ ত্বরণ
৬৪K সিকোয়েন্স দৈর্ঘ্যে মূল MoBA বাস্তবায়নের তুলনায় ৭.৪ গুণ ত্বরণ এবং ৬.১ গুণ মেমরি সাশ্রয়
অভিজ্ঞতামূলক যাচাইকরণ: স্ক্র্যাচ থেকে LLMs প্রশিক্ষণের মাধ্যমে উন্নত MoBA মডেল যাচাই করা, ৭/৮ স্পার্সিটি বজায় রেখে ঘন মনোযোগ বেসলাইনের কর্মক্ষমতা মেলানো
ইনপুট: 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 টি সবচেয়ে প্রাসঙ্গিক ব্লক নির্বাচন করা হয়।
১. টাইলড শীর্ষ-K নির্বাচন (Flash TopK)
তিন-পর্যায়ের পাইপলাইন:
পর্যায় ১: Triton কার্নেল কী ব্লক সেন্ট্রয়েড গণনা করে, ছোট ম্যাট্রিক্স K̃ উৎপন্ন করে
পর্যায় ২: FlashAttention-2 দ্বারা অনুপ্রাণিত টাইলড কার্নেল, Q এবং K̃ এর মধ্যে স্কোর গণনা করে, প্রতিটি প্রশ্নের শীর্ষ-k কী ব্লক খুঁজে পায়, সম্পূর্ণ স্কোর ম্যাট্রিক্স উপাদান ছাড়াই (অ্যালগরিদম ৩)
পর্যায় ৩: উচ্চ দক্ষ epilogue প্রশ্ন কেন্দ্র সূচক কী ব্লক কেন্দ্রের varlen লেআউটে পুনর্বিন্যাস করে
প্রতিটি যৌক্তিক প্রশ্ন ব্লক Q_i এর জন্য:
প্রতিটি যৌক্তিক কী ব্লক K_j এর জন্য:
varlen সূচক ব্যবহার করে প্রাসঙ্গিক প্রশ্ন খুঁজে পান
প্রশ্ন উপসেটকে ঘন ভৌত ব্লকে ব্যাচ করুন:
- HBM থেকে ভৌত প্রশ্ন ব্লক SRAM এ সংগ্রহ করুন
- SRAM এ ক্যাশ করুন, যৌক্তিক কী ব্লক K_j এর সমস্ত ভৌত টাইল জুড়ে পুনরায় ব্যবহার করুন
- উচ্চ দক্ষ ঘন GEMM সম্পাদন করুন
- ফলাফল HBM এ ছড়িয়ে দিন
মূল অপ্টিমাইজেশন: SRAM এ সংগৃহীত প্রশ্ন ব্লক ক্যাশ করে, একাধিক ঘন GEMM জুড়ে পুনরায় ব্যবহার করে, অনিয়মিত সংগ্রহ অপারেশনের খরচ কার্যকরভাবে পরিশোধ করে
३. ব্যাকওয়ার্ড পাস: পুনঃগণনা (অ্যালগরিদম ৫)
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 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) - মিক্সচার অফ ব্লক অ্যাটেনশন ধারণা প্রস্তাব
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 উন্নয়নের জন্য উল্লেখযোগ্য। বিশেষভাবে প্রশংসনীয় হল লেখকদের নিয়ন্ত্রণ ভেরিয়েবল পরীক্ষার মাধ্যমে তত্ত্ব যাচাইকরণের কঠোর মনোভাব এবং সম্প্রদায় গ্রহণ প্রচারের জন্য ওপেন সোর্স কোডের প্রচেষ্টা।