Optimizing Mixture of Block Attention
Not on the map. Why?
Abstract (arXiv)
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.
Key facts. Cited by 8 papers in the map (data through 7 Oct 2026). Builds on 35 of its 41 references that are in the map. First posted on arXiv 14 Nov 2025.
Builds on
35 of 41 references.
References · 35
- 2017Attention Is All You Needcited by 59,932
- 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,767
- 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,710
- 2024Qwen2-VL: Enhancing Vision-Language Model's Perception of the World at Any Resolutioncited by 4,151
- 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,018
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
- 2020Longformer: The Long-Document Transformercited by 3,362
- 2016Pointer Sentinel Mixture Modelscited by 2,992
- 2021TruthfulQA: Measuring How Models Mimic Human Falsehoodscited by 2,919
- 2019PIQA: Reasoning about Physical Commonsense in Natural Languagecited by 2,818
- 2023Efficient Streaming Language Models with Attention Sinkscited by 2,115
- 2018Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answeringcited by 2,114
23 more
- 2020Big Bird: Transformers for Longer Sequencescited by 1,673
- 2020Reformer: The Efficient Transformercited by 1,658
- 2024HunyuanVideo: A Systematic Framework For Large Video Generative Modelscited by 1,597
- 2020Rethinking Attention with Performerscited by 1,584
- 2019Generating Long Sequences with Sparse Transformerscited by 1,523
- 2020Linformer: Self-Attention with Linear Complexitycited by 1,462
- 2023Video-LLaVA: Learning United Visual Representation by Alignment Before Projectioncited by 1,385
- 2023LongBench: A Bilingual, Multitask Benchmark for Long Context Understandingcited by 1,378
- 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,311
- 2024RULER: What's the Real Context Size of Your Long-Context Language Models?cited by 1,270
- 2016The LAMBADA dataset: Word prediction requiring a broad discourse contextcited by 960
- 2023H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Modelscited by 805
- 2019Fast Transformer Decoding: One Write-Head is All You Needcited by 700
- 2024Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inferencecited by 487
- 2020Efficient Content-Based Sparse Attention with Routing Transformerscited by 469
- 2023Ring Attention with Blockwise Transformers for Near-Infinite Contextcited by 448
- 2025Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attentioncited by 391
- 2024Parallelizing Linear Transformers with the Delta Rule over Sequence Lengthcited by 373
- 2025FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Servingcited by 248
- 2019Blockwise Self-Attention for Long Document Understandingcited by 158
- 2025XAttention: Block Sparse Attention with Antidiagonal Scoringcited by 131
- 2023FlashDecoding++: Faster Large Language Model Inference on GPUscited by 70
- 2025SWAN-GPT: An Efficient and Scalable Approach for Long-Context Language Modelingcited by 16
Built on by
Citing papers per year: 2026: 8.
8 papers.
- 2026MiniMax Sparse Attentioncited by 34
- 2026FlashPrefill: Instantaneous Pattern Discovery and Thresholding for Ultra-Fast Long-Context Prefillingcited by 9
- 2026Sol-Attn: Accelerating Video Generation Inference via On-the-Fly Attention Sparsificationcited by 6
- 2026UniPrefill: Universal Long-Context Prefill Acceleration via Block-wise Dynamic Sparsificationcited by 2
- 2026Can Language Models Actually Retrieve In-Context? Drowning in Documents at Million Token Scalecited by 2
- 2026FlashPrefill V2: Block-Sparse Prefill Attention for Long-Context LLM Servingcited by 2
- 2026Power-based Partial Attention: Bridging Linear-Complexity and Full Attentioncited by 0
- 2026Superlinear Multi-Step Attentioncited by 0
- 2026FlashPrefill V2: Block-Sparse Prefill Attention for Long-Context LLM Serving20 Aug 2026
- 2026Sol-Attn: Accelerating Video Generation Inference via On-the-Fly Attention Sparsification27 Jul 2026
- 2026Can Language Models Actually Retrieve In-Context? Drowning in Documents at Million Token Scale1 Jul 2026
- 2026MiniMax Sparse Attention11 Jun 2026
- 2026UniPrefill: Universal Long-Context Prefill Acceleration via Block-wise Dynamic Sparsification7 May 2026
- 2026FlashPrefill: Instantaneous Pattern Discovery and Thresholding for Ultra-Fast Long-Context Prefilling6 Mar 2026
- 2026Superlinear Multi-Step Attention26 Jan 2026
- 2026Power-based Partial Attention: Bridging Linear-Complexity and Full Attention24 Jan 2026