MassAlloc Attention: Let Attention Allocate Its Own Compute
Not on the map. Why?
Abstract (arXiv)
FullAttn often assigns negligible normalized mass to much of the causal score space, yet dense kernels execute the complete post-score path after forming each QK tile. We introduce MALA, a fused attention primitive that preserves score access to every legal causal interaction and uses normalized contribution to allocate post-score computation. Forward uses its evolving online-softmax normalizer, while backward reuses the finalized normalizer to derive nested retained support using only standard attention state. A common tolerance governs training and inference, allowing for adaptive retention of the work. MALA reduces low-contribution post-score computation. A matched-work study at 8K isolates the benefit of distribution-adaptive allocation: under exactly matched total post-score work, MALA approaches a per-instance reference-mass oracle, with mean omitted mass of 0.0188% versus 0.0182%. Across context lengths from 1K to 32K tokens, the same tolerance maintains low output and gradient errors relative to the reference. Across a broader controlled associative-recall comparison, MALA closely tracks FullAttn as context grows, reaching 89.67% accuracy at 8K compared with 89.97% for FullAttn. In an attention-operator benchmark at 128K tokens with tensor parallelism, MALA reduces forward and backward latency during training by 2.2x and 3.0x and decoding latency during inference by 1.6x relative to FullAttn. Across scaling-law training from 0.6B to 14B parameters, MALA closely tracks FullAttn in perplexity while reducing total training FLOPs. The resulting 14B models and 32B models from separate continued training achieve comparable knowledge, reasoning, and long-context retrieval scores to FullAttn. These results indicate that allocating post-score computation according to normalized attention contributions can retain the evaluated capabilities of FullAttn while reducing attention computation.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 51 of its 66 references that are in the map. First posted on arXiv 26 Sep 2026.
Builds on
51 of 66 references.
References · 51
- 2017Attention Is All You Needcited by 60,361
- 2021Training Verifiers to Solve Math Word Problemscited by 10,109
- 2020Measuring Massive Multitask Language Understandingcited by 7,963
- 2020Scaling Laws for Neural Language Modelscited by 6,586
- 2021Measuring Mathematical Problem Solving With the MATH Datasetcited by 5,371
- 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,886
- 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,100
- 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
- 2023Generative Agents: Interactive Simulacra of Human Behaviorcited by 3,568
- 2020Longformer: The Long-Document Transformercited by 3,413
- 2019PIQA: Reasoning about Physical Commonsense in Natural Languagecited by 2,883
- 2023GPQA: A Graduate-Level Google-Proof Q&A Benchmarkcited by 2,879
38 more
- 2018Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answeringcited by 2,164
- 2024Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameterscited by 1,986
- 2019Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelismcited by 1,868
- 2017Decoupled Weight Decay Regularizationcited by 1,831
- 2022Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Themcited by 1,807
- 2020Big Bird: Transformers for Longer Sequencescited by 1,684
- 2019Generating Long Sequences with Sparse Transformerscited by 1,546
- 2024RULER: What's the Real Context Size of Your Long-Context Language Models?cited by 1,346
- 2024SnapKV: LLM Knows What You are Looking for Before Generationcited by 880
- 2023H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Modelscited by 832
- 2023AGIEval: A Human-Centric Benchmark for Evaluating Foundation Modelscited by 716
- 2023YaRN: Efficient Context Window Extension of Large Language Modelscited by 601
- 2024FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precisioncited by 519
- 2024Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inferencecited by 501
- 2025Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attentioncited by 407
- 2023Effective Long-Context Scaling of Foundation Modelscited by 303
- 2024DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Headscited by 288
- 2024CodeAgent: Enhancing Code Generation with Tool-Integrated Agent Systems for Real-World Repo-level Coding Challengescited by 282
- 2025FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Servingcited by 258
- 2024When Attention Sink Emerges in Language Models: An Empirical Viewcited by 208
- 2023Zoology: Measuring and Improving Recall in Efficient Language Modelscited by 204
- 2025MoBA: Mixture of Block Attention for Long-Context LLMscited by 198
- 2024InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memorycited by 197
- 2024Flex Attention: A Programming Model for Generating Optimized Attention Kernelscited by 185
- 2024Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durationscited by 170
- 2018Online normalizer calculation for softmaxcited by 162
- 2025Why do LLMs attend to the first token?cited by 127
- 2024SeerAttention: Learning Intrinsic Sparse Attention in Your LLMscited by 124
- 2025FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inferencecited by 123
- 2025SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inferencecited by 120
- 2024ClusterKV: Manipulating LLM KV Cache in Semantic Space for Recallable Compressioncited by 65
- 2025Attention Sinks and Compression Valleys in LLMs are Two Sides of the Same Coincited by 52
- 2025Twilight: Adaptive Attention Sparsity with Hierarchical Top-p Pruningcited by 49
- 2025InfLLM-V2: Dense-Sparse Switchable Attention for Seamless Short-to-Long Adaptationcited by 32
- 2025Sliding Window Attention Training for Efficient Large Language Modelscited by 24
- 2024Beyond Probabilities: Unveiling the Misalignment in Evaluating Large Language Modelscited by 19
- 2025BLASST: Dynamic BLocked Attention Sparsity via Softmax Thresholdingcited by 13
- 2026Double-P: Hierarchical Top-P Sparse Attention for Long-Context LLMscited by 5
50 most cited of 51 All 51 (JSON)
Built on by
Not cited yet.