MassAlloc Attention: Let Attention Allocate Its Own Compute
cs.AI, cs.CL
Submitted: 2026-09-26
Updated: 2026-09-26
Code: https://github.com/huggingface/open-r1
License: http://creativecommons.org/licenses/by/4.0/
The gist: 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.
Terminology
Abstract
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.
Sources
- Why do LLMs attend to the first token?
- Longformer: The Long-Document Transformer
- Generating Long Sequences with Sparse Transformers
- Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge
- Training Verifiers to Solve Math Word Problems
- Attention Sinks and Compression Valleys in LLMs are Two Sides of the Same Coin
- DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models
- Flex Attention: A Programming Model for Generating Optimized Attention Kernels
- Sliding Window Attention Training for Efficient Large Language Models
- SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs
- When Attention Sink Emerges in Language Models: An Empirical View
- DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning
- Measuring Mathematical Problem Solving With the MATH Dataset
- RULER: What's the Real Context Size of Your Long-Context Language Models?
- NOSA: Native and Offloadable Sparse Attention
- Scaling Laws for Neural Language Models
- FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference
- Predictable Scale: Part I, Step Law -- Optimal Hyperparameter Scaling Law in Large Language Model Pretraining
- Twilight: Adaptive Attention Sparsity with Hierarchical Top-$p$ Pruning
- ClusterKV: Manipulating LLM KV Cache in Semantic Space for Recallable Compression
Related papers
- MAVEN-T: Reinforced Heterogeneous Distillation for Real-Time Multi-Agent Trajectory Prediction
- Model Discovery Agent: LLM-assisted Bayesian experiment design for data-efficient discovery of mechanistic world models
- Agents' Last Exam
- The Clinician's Veto: Navigating Trust, Liability, and Uncertainty in Autonomous AI Prescribing
- MindHelper: Closed-Loop Embodied Mental-State Reasoning for Precision Intervention
- Incumbent Advantage: Brand Bias and Cognitive Manipulation Dynamics in LLM Recommendation Systems