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
cs.AI, cs.CL
2026-09-26
MALA keeps full QK scoring and skips low-mass softmax/PV. 14B 32K training FLOPs drop 23.1%; at 128K, backward is 3.0× faster than FullAttn with matched eval scores.
FlashAttention tiles and fuses QK, softmax, and PV so the operator hits memory less often. The execution pattern stays dense: after every legal causal tile forms its scores, the kernel still runs the full post-score path, softmax updates, V loads, PV accumulation, and on the backward pass probability reconstruction plus dP, dS, dQ, dK, and dV. On long contexts the normalized mass piles up on local windows, sink tokens, and a sparse retrieval set. Many remaining tiles get almost nothing, and the dense kernel still finishes the post-score work.
Static windows and block patterns prescribe support by position. Dynamic sparse methods drop interactions before exact attention, so they can skip QK too. MALA splits the loop: every legal causal tile still gets QK, and only post-score work is gated by normalized contribution. The open question is whether attention's own scores and softmax statistics can allocate that work without losing operator fidelity, language-model quality, or long-range retrieval.
MassAlloc Attention (MALA) is a fused primitive from HKUST-GZ and BAAI. After each QK tile, it compares the largest per-key contribution ratio in the tile against a length-normalized threshold τ/Lq. Lq is the number of causally visible keys for that query; uniform attention would give each key 1/Lq, and τ scales that reference. Sharp rows can drop many tiles. Diffuse rows keep more. The default τ=1 is shared across training forward and backward, prefill, and decoding. Realized work then varies by query, head, layer, and input.
Forward uses the evolving online-softmax normalizer. The test is conservative at tile granularity: one row that fails the skip keeps the tile for the whole query block. Skipped tiles do not load V or update softmax and PV. Tiles are visited from the diagonal toward earlier keys so the running normalizer is established on recent context. That order is an execution prior, not a support assumption; a distant tile with a large enough score is kept. The first legal tile is always retained, so no row is empty. Forward writes the final log-normalizer into ordinary attention state.
Backward tests against that finalized normalizer rather than replaying partial sums. Because the final denominator is larger, contribution ratios can only shrink, so every forward-skipped tile is skipped in backward as well. Backward may drop extra tiles that forward kept before the normalizer was complete. Nested support needs no stored mask. Backward is therefore not an exact derivative of the forward operator; gradient error is measured, not proved.
QK stays quadratic. Savings come from skipped post-score arithmetic and traffic. As τ goes to 0, both passes recover FullAttn. The kernel materializes neither the attention matrix nor a selection mask. The persistent KV cache is neither compressed nor evicted.
At 8K, every policy in the matched-work study runs the same total post-score work, about 1,024 key slots per query. Position-only static allocation omits 0.6012% mean probability mass. Static layer-head-position allocation omits 0.1721%. A per-instance reference-mass oracle omits 0.0182%. MALA, using only online state, omits 0.0188% and has 0.0174% mean relative output L2 error against the oracle's 0.0164%.
With the same τ=1 from 1K to 32K, mean forward post-score work rises from 470 to 1,066 slots per query and stays between 1,010 and 1,066 past 8K. Mean omitted mass is at most 0.0062%; mean relative output L2 error is at most 0.021%. Mean relative gradient errors stay within 0.38% (dQ), 0.35% (dK), and 0.17% (dV).
On controlled associative recall with 256 bindings and no distractors, fixed-budget sparse methods share a 1,024-slot ceiling. At 8K and dmodel=512:
| Method | Accuracy |
| FullAttn | 89.97% |
| MALA | 89.67% |
| DSA | 52.61% |
| MoBA | 47.25% |
| NSA | 22.61% |
At 128K on 8 H100s with tensor parallelism TP=8, MALA cuts training forward and backward latency by 2.2× and 3.0× and decoding latency by 1.6× versus FullAttn, at FullAttn-level peak operator memory. Decode latency is within 3% of DSA.
From 0.6B to 14B, 14B perplexity differs from FullAttn by less than 0.001. Total training FLOPs fall 2.5% in 4K pre-training and 23.1% in 32K long-context training. The 14B models are trained from scratch. The 32B pair continues from Qwen3-32B for 64B tokens at 32K.
| Scale | Method | Knowledge | Reasoning | RULER 32K | RULER 128K |
| 14B | FullAttn | 72.32 | 64.46 | 89.42 | 65.84 |
| 14B | MALA | 72.48 | 64.66 | 89.45 | 65.75 |
| 14B | MoBA | 70.01 | 58.07 | 86.85 | 60.17 |
| 14B | DSA | 70.11 | 62.04 | 87.59 | 63.05 |
| 32B | FullAttn | 75.62 | 75.67 | 92.70 | 82.03 |
| 32B | MALA | 75.75 | 76.10 | 92.71 | 82.56 |
The 23.1% FLOP cut at 14B and 32K is whole-run training compute, not a kernel-only number. The 3.0× backward speedup at 128K is the more useful training lever. Decode is only 1.6× faster because QK discovery remains quadratic and post-score work is a smaller share of inference.
MoBA, DSA, and NSA select support before scoring. MALA does not drop interactions before QK, which is why associative recall stays next to FullAttn while fixed-budget methods lose most bindings. QK cannot be skipped. The 2.5% saving at 4K pre-training is small; the method pays off on long sequences.
Code is released as flash-sparse-attention. One tolerance covers train and serve, with no per-layer, per-head threshold calibration. The fit is teams already on FlashAttention-style fused kernels who are pushing 32K to 128K and do not want a new sparse topology. This is a constant-factor execution change, not a new attention function.
QK is still O(n²). The paper is explicit: support-sparse methods can skip QK on excluded regions; MALA cannot. A 2.5% FLOP cut at 4K pre-training will not rewrite short-context scaling.
Backward support is nested in forward support, but it is not an exact adjoint. Extra omitted gradient terms depend on dO and V; fidelity is empirical. The omitted-mass bound is τ/(1+τ), which is 50% at τ=1, while measured omitted mass is a few thousandths of a percent. The bound is loose. Tile tests are conservative: one live row keeps the whole query block.
The KV cache is left intact, so serving memory does not shrink, only the operator working set. A 1.6× decode speedup can be dominated by cache traffic. The 32B models are continued training from Qwen3-32B, not from-scratch MALA runs. Associative recall has no distractors; it tests whether 256 bindings survive, not needle-finding in noise. The operator benchmark excludes QKV projections, MLPs, the optimizer, and pipeline communication, so end-to-end step speedups will be smaller than 2.2× and 3.0×.