mlsanityData to 7 Oct 2026

Optimizing Mixture of Block Attention

arXiv 2511.11571 ↗First posted 14 Nov 2025cs.LG cs.CLcited by 8

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

  1. 2017Attention Is All You Needcited by 59,932
  2. 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,767
  3. 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,710
  4. 2024Qwen2-VL: Enhancing Vision-Language Model's Perception of the World at Any Resolutioncited by 4,151
  5. 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,018
  6. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
  7. 2020Longformer: The Long-Document Transformercited by 3,362
  8. 2016Pointer Sentinel Mixture Modelscited by 2,992
  9. 2021TruthfulQA: Measuring How Models Mimic Human Falsehoodscited by 2,919
  10. 2019PIQA: Reasoning about Physical Commonsense in Natural Languagecited by 2,818
  11. 2023Efficient Streaming Language Models with Attention Sinkscited by 2,115
  12. 2018Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answeringcited by 2,114
23 more
  1. 2020Big Bird: Transformers for Longer Sequencescited by 1,673
  2. 2020Reformer: The Efficient Transformercited by 1,658
  3. 2024HunyuanVideo: A Systematic Framework For Large Video Generative Modelscited by 1,597
  4. 2020Rethinking Attention with Performerscited by 1,584
  5. 2019Generating Long Sequences with Sparse Transformerscited by 1,523
  6. 2020Linformer: Self-Attention with Linear Complexitycited by 1,462
  7. 2023Video-LLaVA: Learning United Visual Representation by Alignment Before Projectioncited by 1,385
  8. 2023LongBench: A Bilingual, Multitask Benchmark for Long Context Understandingcited by 1,378
  9. 2023GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpointscited by 1,311
  10. 2024RULER: What's the Real Context Size of Your Long-Context Language Models?cited by 1,270
  11. 2016The LAMBADA dataset: Word prediction requiring a broad discourse contextcited by 960
  12. 2023H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Modelscited by 805
  13. 2019Fast Transformer Decoding: One Write-Head is All You Needcited by 700
  14. 2024Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inferencecited by 487
  15. 2020Efficient Content-Based Sparse Attention with Routing Transformerscited by 469
  16. 2023Ring Attention with Blockwise Transformers for Near-Infinite Contextcited by 448
  17. 2025Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attentioncited by 391
  18. 2024Parallelizing Linear Transformers with the Delta Rule over Sequence Lengthcited by 373
  19. 2025FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Servingcited by 248
  20. 2019Blockwise Self-Attention for Long Document Understandingcited by 158
  21. 2025XAttention: Block Sparse Attention with Antidiagonal Scoringcited by 131
  22. 2023FlashDecoding++: Faster Large Language Model Inference on GPUscited by 70
  23. 2025SWAN-GPT: An Efficient and Scalable Approach for Long-Context Language Modelingcited by 16

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 8.

8 papers.