PaperlayerView exact revision on arXiv

Source-linked AI summary

MassAlloc Attention: Let Attention Allocate Its Own Compute

Jingze Shi, Zhangyang Peng, Xianduo Li, Yanlin Qi, Xiaotian Lin, Haoxian Chen, Liangdong Wang, Guang Liu, Yuyu Luo

arXiv:2609.32712v1cs.AIcs.CL

TL;DR

Dense attention performs post-score computation even for interactions with negligible normalized mass. MALA preserves access to every legal causal score and allocates post-score work from normalized contributions. Across the reported evaluations, it closely tracks FullAttn’s quality while reducing attention computation.

  • Problem

    Dense kernels execute the full post-score path for every legal causal tile, including tiles assigned negligible normalized mass.

    Source: §1 · pp. 1–2

  • Method

    MALA preserves causal QK score discovery and uses normalized contributions with a shared tolerance to allocate post-score execution in forward and backward.

    Source: §1 · pp. 2–3

  • Results

    MALA closely tracks FullAttn in associative recall, perplexity, and model-level evaluations while reducing attention latency and training FLOPs.

    Source: §1 · p. 3, §5 · p. 9

  • Takeaways & Limitations

    The evaluated results indicate that normalized-contribution allocation can retain FullAttn-level capabilities while reducing attention computation.

    Source: §5 · p. 9

    Backward can omit additional gradient contributions, so its nested retained support does not exactly differentiate the forward operator.

    Source: §2.2 · p. 5

Abstract

from arXiv · show

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.

1 INTRODUCTION

MALA allocates post-score computation using normalized attention contributions while preserving score access to every legal causal interaction. Across the reported evaluations, it closely tracks FullAttn’s quality while reducing attention computation.

1 INTRODUCTION

  • Dense attention executes post-score work for every legal causal tile, including tiles with negligible normalized mass.That work includes softmax updates, V loading, and PV accumulation in forward, plus probability reconstruction and gradient computations in backward.Source: §1 · pp. 1–2
  • MALA uses normalized contribution to allocate post-score work, with online-normalizer tests in forward and finalized-normalizer tests in backward.A common tolerance governs training and inference, and backward support is nested within forward support.Source: §1 · p. 2
  • MALA preserves full causal QK score discovery, so its savings come from skipping low-contribution post-score execution rather than reducing QK complexity.Its forward rule provides an omitted-mass guarantee, while output and gradient fidelity are evaluated empirically.Source: §1 · p. 2
  • At 8K, mean omitted mass is 0.0188% for MALA versus 0.0182% for the per-instance reference-mass oracle under exactly matched post-score work.Static position- and layer-head-based allocations perform substantially worse.Source: §1 · p. 2
  • Across evaluations, MALA closely tracks FullAttn in recall, perplexity, and model-level scores while reducing attention latency and training FLOPs.Model-level comparisons cover knowledge, reasoning, and long-context retrieval at 14B and 32B parameters.Source: §1 · p. 3
  • The paper frames sparse attention as distribution-conditioned compute allocation and evaluates it through matched-work controls, fidelity, recall, kernel, scaling, and model-level studies.The fused operators use native attention state and a common normalized-mass tolerance.Source: §1 · p. 3

2 METHODOLOGY

MALA separates QK score discovery from post-score execution, using a length-aware contribution tolerance to decide whether each tile proceeds through the remaining attention computation. Its forward and backward rules share standard attention state, and its savings reduce post-score work while retaining quadratic QK computation.

2 METHODOLOGY

  • The methodology treats attention as a runtime compute allocator, separating score discovery from post-score execution and pairing online forward with offline backward allocation.The method uses standard attention state rather than a stored forward mask.Source: §2 · p. 3, §2.1 · p. 3, §2.4 · p. 5

2.1 ATTENTION AS A COMPUTE ALLOCATOR

  • MALA allocates post-score computation after forming QK scores for every legal causal tile, with a length-aware threshold τ/Lq based on uniform-attention mass.The tolerance specifies admissible contribution rather than a fixed compute budget; sharper rows can reject more tiles, while diffuse rows generally retain more.Source: §2.1 · pp. 3–4

2.2 ONLINE AND OFFLINE ALLOCATION

  • Forward tests tiles using an evolving online-softmax normalizer, while backward tests recomputed tiles with the finalized normalizer and derives support nested within forward support.Backward may omit additional gradient contributions, so nested support does not imply exact differentiation of the forward operator.Source: §2.2 · pp. 4–5

2.3 ATTENTION COST

  • MALA remains quadratic in sequence length because it computes every legal causal QK tile; savings are data-dependent reductions in post-score arithmetic and memory traffic.Forward skips avoid exponentiation, V loading, and PV accumulation; backward skips avoid probability reconstruction and gradient-path operations.Source: §2.3 · p. 5

2.4 FUSED EXECUTION

  • The fused operators keep attention state on chip and materialize neither the attention matrix nor a mask, retained-tile indices, or router state.Split-KV decoding tests tiles against each split’s partial normalizer; MALA does not compress or evict the persistent KV cache.Source: §2.4 · p. 5

3 EXPERIMENTS

The experiments test whether normalized-mass allocation preserves operator behavior, retrieval, and language-model quality while reducing attention computation. They evaluate matched-work allocation, operator fidelity, latency, scaling, and model-level capabilities.

Normalized-Mass Allocation at Matched Work

  • With exactly matched post-score work of approximately 1,024 slots per query, MALA’s mean omitted mass is 0.0188% versus 0.0182% for the per-instance reference-mass oracle.The comparison isolates distribution-adaptive allocation; MALA’s mean relative output L2 error is 0.0174% versus 0.0164% for the oracle.Source: Table 1 · p. 6, §3 · p. 6

Operator Fidelity Across Context Lengths

  • Across 1K–32K contexts, mean omitted mass is at most 0.0062%, mean relative output L2 error at most 0.021%, and mean relative gradient errors at most 0.38% for dQ, 0.35% for dK, and 0.17% for dV.Mean forward and backward work rises from 470 and 462 key slots per query to 1,066 and 1,058 as context grows 32×.Source: Figure 2 · p. 6, §3 · pp. 6–7

Controlled Associative Recall

  • At 8,192 tokens and dmodel = 512, MALA reaches 89.67% accuracy versus 89.97% for FullAttn; DSA, MoBA, and NSA reach 52.61%, 47.25%, and 22.61%.MALA closely tracks FullAttn as context grows under a 1,024-slot ceiling for fixed-budget methods.Source: Figure 3 · p. 7, §3 · p. 7

Operator Latency and Memory

  • At 128K, MALA reduces training forward and backward latency by 2.2× and 3.0×, and inference decoding latency by 1.6× versus FullAttn.Its peak allocation matches FullAttn across the three operators and measured context lengths.Source: Figure 4 · p. 7, §3 · p. 8

Scaling Across Model Sizes

  • At 14B, MALA’s Perplexity differs from FullAttn by less than 0.001 after both training stages, while total FLOPs fall by 2.5% and 23.1%.Across five model sizes from 0.6B to 14B, MALA closely tracks FullAttn while shifting the quality-compute frontier toward lower FLOPs.Source: Figure 5 · p. 8, §3 · p. 8

Model-Level Evaluation

  • At 14B and 32B, MALA remains comparable to FullAttn in knowledge, reasoning, native 32K retrieval, and YaRN-extrapolated 128K retrieval.For example, at 14B, knowledge and reasoning scores are 72.48 and 64.66 for MALA versus 72.32 and 64.46 for FullAttn.Source: Table 2 · p. 8, §3 · pp. 8–9

4 RELATED WORK

MALA preserves complete causal QK score discovery and allocates post-score computation using scores and softmax state. The related work contrasts this design with support restriction, adaptive allocation, hardware-aware execution, and KV-cache management.

Sparse attention

  • Unlike sparse methods that restrict support before exact attention, MALA discovers scores over the complete causal context and allocates post-score work from those scores and softmax state.Fixed sparse attention uses local windows, structured blocks, or global tokens [3] [7] [61]; dynamic methods select tokens or blocks using scores, estimates, cache statistics, or routers [13] [17] [30] [34] [49] [55] [60].Source: §4 · p. 9

Adaptive allocation and threshold calibration

  • MALA applies one normalized-contribution tolerance across heads, layers, and inputs, using exact QK scores and native softmax normalizers rather than separately calibrated thresholds.Twilight and Double-P use hierarchical top-p selection [31] [38], while Top-Theta calibrates thresholds offline.Source: §4 · p. 9

Hardware-aware sparse execution

  • MALA tests scores against accumulated or finalized softmax normalizers, supporting an omitted-mass guarantee and fused forward and backward operators with nested backward support.FlashAttention uses tiled execution and online softmax [10] [37] [45]; SpargeAttention combines prediction with an online filter [63], while BLASST calibrates length-dependent thresholds [59].Source: §4 · p. 9

KV-cache management

  • MALA retains the full KV history for score discovery and value access, leaving cache storage and placement as separate optimization opportunities.KV-cache systems reduce inference cost by pruning, compressing, retrieving, or offloading cached states [25] [30] [32] [49] [55].Source: §4 · p. 9

5 CONCLUSION

MALA allocates post-score computation from normalized attention contributions while preserving complete causal QK score discovery, and the conclusion reports comparable evaluated capabilities with reduced attention computation. Its fused forward, backward, and decoding procedures use standard attention state to make allocation decisions.

5 CONCLUSION

  • MALA preserves complete causal QK score discovery; its savings come from skipping low-contribution post-score execution under a shared tolerance.Forward uses online allocation, while backward uses the saved normalizer to derive nested retained support from standard attention state.Source: §5 · p. 9
  • MALA’s evaluated models retain comparable FullAttn-level capabilities while reducing attention computation, with lower latency and training FLOPs reported across the evaluations.The conclusion reports comparable perplexity, knowledge, reasoning, and long-context retrieval, alongside lower attention-operator latency and training FLOPs.Source: §5 · p. 9

A COMPLETE FORWARD AND BACKWARD ALGORITHMS

  • The complete tiled algorithms test candidate tiles before the forward value/softmax path and recompute QK before backward decides whether to execute retained-tile work.Forward maintains the online-softmax state and stores the final normalizer; backward reconstructs probabilities for retained tiles from that saved state.Source: A COMPLETE FORWARD AND BACKWARD… · p. 16, §A.1 · p. 17

A.1 FUSED EXECUTION DETAILS

  • Fused forward, backward, and decoding embed allocation in tiled attention loops, with one probability tolerance and execution-specific realized work.Decoding evaluates causal QK tiles while retaining the KV cache; split-KV partial outputs are combined by the standard online-softmax reduction.Source: §A.1 · pp. 16–17

B FORWARD APPROXIMATION GUARANTEES

The forward rule bounds the dense-softmax probability mass assigned to omitted keys. The formal guarantee assumes exact arithmetic; the fused implementation’s finite-precision bounds hold up to rounding error.

B FORWARD APPROXIMATION GUARANTEES

  • The omitted dense-softmax mass is bounded by ϵ_q ≤ τ/(1 + τ) for a query with L_q causally visible keys.The bound follows from the forward rule in Equation 3; the passage also defines the omitted-mass fraction as ϵ_q = (Z_O/Z_R)/(1 + Z_O/Z_R).Source: B FORWARD APPROXIMATION GUARANTEES · p. 17
  • The formal statements assume exact arithmetic, while fused finite-precision tests satisfy the bounds up to corresponding rounding error.The implemented operators’ omitted mass and output and gradient fidelity are measured empirically.Source: B FORWARD APPROXIMATION GUARANTEES · p. 18

C EXPERIMENT SETTINGS

The experiments use matched model and training settings while distinguishing realized post-score work, FLOP-equivalent work, and full operator cost. Their setups compare adaptive allocation with matched-work controls and measure operator fidelity across context lengths.

C EXPERIMENT SETTINGS

  • Experiments match model scale, depth, hidden size, data, and optimization settings, with MALA generally using τ = 1 and separate H100 configurations for training and benchmarks.Scaling-law and continued-training experiments use 128 H100 GPUs; operator benchmarking uses 8 H100 GPUs with tensor parallelism TP = 8.Source: C EXPERIMENT SETTINGS · p. 18

C.1 WORK AND COST ACCOUNTING

  • Work accounting separates realized post-score slots, architecture-normalized FLOP-equivalent work, and full operator cost including score discovery and selection.Raw slots support within-architecture and associative-recall comparisons; normalized work is used for cross-architecture operator and scaling-law comparisons.Source: §C.1 · pp. 18–19

C.2 NORMALIZED-MASS ALLOCATION AT MATCHED WORK SETUP

  • The matched-work study assigns all policies the same total post-score slots, then contrasts static position- or layer-head-position allocations with per-instance allocation.Controls use the screening-disabled reference to select highest-mass regions; the per-instance control uses MALA’s realized slot count for each input, position, layer, and head.Source: §C.2 · p. 19

C.3 OPERATOR FIDELITY ACROSS CONTEXT LENGTHS SETUP

  • Operator fidelity is measured on 256 sequences with prefixes from 1,024 to 32,768 tokens, using relative L2 errors for output and dQ, dK, and dV.Backward uses a deterministic Rademacher probe for dO; measurements aggregate complete operators at head level and bootstrap confidence intervals by sequence.Source: §C.3 · p. 19

C.4 ASSOCIATIVE RECALL SETUP

The associative-recall setup compares methods at a common raw-slot scale, while the operator benchmark measures complete attention paths across sequence lengths on tensor-parallel H100 GPUs. The two setups control different quantities: retained interactions in recall and reference-equivalent post-score work in the operator benchmark.

C.4 ASSOCIATIVE RECALL SETUP

  • Associative recall uses 256 key-value pairs across lengths of 1,024–8,192 tokens, with fixed-budget sparse methods capped at 1,024 post-score slots and MALA using τ = 1.The raw-slot comparison aligns retained interactions; DSA shares the ceiling despite its larger value dimension.Source: §C.4 · p. 20, Table 3 · p. 20

C.5 OPERATOR LATENCY AND MEMORY SETUP

  • The operator benchmark spans 1,024–131,072 tokens on 8 H100 GPUs with TP = 8, targeting comparable reference-equivalent forward post-score work for sparse methods.Latency and memory include score-discovery and selection paths; the benchmark excludes model-level projections, MLPs, optimizer updates, and pipeline communication.Source: §C.5 · p. 20, Table 4 · p. 21

C.6 SCALING LAWS SETUP

The scaling-law experiments compare attention variants across pre-training and long-context training, with configurations and compute accounting designed to describe their work on comparable terms.

C.6 SCALING LAWS SETUP

  • All scaling-law runs use SmolLMCorpus, the Qwen3 tokenizer, AdamW, a WSD schedule, and random seed 42; models train first at 4K and then at 32K tokens.Long-context training initializes from the final pre-training checkpoint.Source: §C.6 · p. 21
  • The compared variants are FullAttn, MoBA, DSA, and MALA, each retaining its native head configuration and using specified attention budgets or tolerance.MALA uses eight KV heads and τ = 1; its realized mean work is near the 1,024-slot reference level.Source: §C.6 · pp. 21–22
  • Sparse variants approximately match forward post-score FLOP-equivalent work across their configurations, while FullAttn attends to the complete context as the dense reference.The per-slot post-score factor is 2hdv, and MALA’s allocation control is a normalized-mass tolerance.Source: §C.6 · pp. 21–22
  • The FLOP-axis counts cumulative full-model training computation, while attention-specific accounting includes score discovery, backward QK recomputation, and retained post-score paths.Routing and indexing costs are included for MoBA and DSA, respectively.Source: §C.6 · p. 22, Table 5 · p. 22

C.7 SPARSE ADAPTATION VIA CONTINUED TRAINING

The 32B continued-training comparison starts both models from the same checkpoint and holds the long-context training setup fixed.

C.7 SPARSE ADAPTATION VIA CONTINUED TRAINING

  • Both 32B models start from the same Qwen3-32B checkpoint and train for 64B tokens at a sequence length of 32,768.The batch size is 4M tokens and the constant learning rate is 1 × 10^-5.Source: §C.7 · p. 22
  • FullAttn and MALA use the same 64-layer, 5,120-hidden-size model configuration, with 64 query-output heads and eight KV heads.Both per-head dimensions dqk and dv are 128.Source: §C.7 · p. 22
  • MALA uses the canonical tolerance τ = 1 while retaining complete causal QK score discovery and distribution-determined post-score work.FullAttn attends to the complete causal context.Source: §C.7 · p. 22

C.8 BENCHMARK EVALUATION

The evaluation compares scaling-law models at 14B and continued-training models at 32B using matched benchmark procedures within each scale.

C.8 BENCHMARK EVALUATION

  • The 14B evaluation compares FullAttn, MoBA, DSA, and MALA, while the 32B evaluation compares FullAttn with MALA.The 14B models come from the scaling-law study, and the 32B models from continued training.Source: §C.8 · p. 23
  • Each task’s reported mean and standard deviation aggregate five evaluation runs using seeds 0, 42, 233, 666, and 1234.Prompts, shot counts, decoding settings, seeds, and metric aggregation are held fixed across variants within each model scale.Source: §C.8 · p. 23
  • The benchmark suite covers knowledge, reasoning, and retrieval, with RULER evaluated at native 32K and after YaRN extrapolation to 128K.The named suites include MMLU, BBH, GSM8K, and RULER.Source: §C.8 · p. 23

D DETAILED MODEL-LEVEL RESULTS

MALA’s model-level evaluations report results comparable to FullAttn across knowledge, reasoning, and long-context retrieval at 14B and 32B, including retrieval after 128K extrapolation.

D DETAILED MODEL-LEVEL RESULTS

Loading 2609.32712v1…