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：保留全部QK打分，按归一化质量跳过低贡献后半段计算。14B长训总FLOPs降23.1%，128K训练反向快3.0倍，知识推理贴近FullAttn。
FlashAttention把QK、softmax、PV融进tiling，IO更省，执行模式没变：每个合法因果tile算出分数后，仍走完整后半段，包括softmax更新、加载V、累加PV，反向还要重建概率并算dP、dS、dQ、dK、dV。长上下文里归一化质量常堆在局部、sink token和少数检索位置，大量tile贡献接近零，dense kernel照样算完。
固定窗口按位置规定support，动态稀疏用分数或router先丢掉交互，QK都可以跳。MALA保留全部合法因果QK，只对后半段按归一化贡献做取舍。这篇要验证的是，用注意力自己的分数和softmax统计量来分配计算，算子保真、模型质量和长程检索能不能保住，训练和推理成本能不能降下来。
MALA（MassAlloc Attention）是港科广与北京智源的fused attention primitive。每个合法因果tile先打QK。是否继续，看该tile最大per-key贡献比是否低于长度归一化阈值τ/Lq。Lq是该query可见key数，均匀注意力每个key分到1/Lq，再乘无量纲容差τ。τ管可接受贡献，不管死预算：尖锐行可丢掉很多tile，弥散行会多留。默认τ=1，训练前向、反向、prefill、decode共用，实际工作量随query、头、层和输入变化。
前向用正在演化的online-softmax归一化因子。当前tile行最大值的指数除以已保留质量，得到贡献比；所有有效行都低于阈值才跳过。跳过则不加载V、不更新softmax、不累加PV。遍历从对角线走向更早的key，先用近处立住归一化因子。这是执行顺序，不是support假设，远处分数够大照样保留。第一块合法tile必留，保证每行非空。前向结束把最终log-normalizer写入标准状态，反向直接用。
反向不再重放前向的部分归一化序列。最终分母更大，贡献比更小，前向跳过的tile反向也会跳；反向还能再跳一些前向在归一化未完成时留下的tile。嵌套support不另存mask。反向因此不是前向的精确微分，梯度误差用实验量。
QK全做，复杂度仍二次。省的是低贡献tile的后半段算术和访存。τ趋向0时退回FullAttn。不物化注意力矩阵，也不存选择mask。KV cache不压缩、不驱逐。
8K上所有策略后半段总工作量锁死，约每query 1024个key slot。按位置静态分配，平均漏掉的概率质量0.6012%；按层-头-位置静态分配0.1721%；逐实例参考质量oracle是0.0182%。MALA只用在线状态，漏质量0.0188%，相对输出L2误差0.0174%，对上oracle的0.0164%。
同一τ=1，上下文从1K到32K。前向后半段工作从每query 470个slot升到1066，8K之后稳定在1010至1066。平均漏质量不超过0.0062%，平均输出相对L2误差不超过0.021%；dQ、dK、dV平均相对误差不超过0.38%、0.35%、0.17%。
受控associative recall用256对绑定、无distractor，固定预算方法天花板1024 slot。8K且dmodel=512时：
| 方法 | 准确率 |
| FullAttn | 89.97% |
| MALA | 89.67% |
| DSA | 52.61% |
| MoBA | 47.25% |
| NSA | 22.61% |
128K、8张H100、张量并行TP=8，相对FullAttn训练前向快2.2倍、反向快3.0倍、decode快1.6倍，峰值算子显存同级；decode延迟与DSA相差不超过3%。0.6B到14B的scaling里，14B perplexity与FullAttn相差不到0.001；总训练FLOPs在4K预训练降2.5%，32K长训降23.1%。
| 规模 | 方法 | 知识均分 | 推理均分 | 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 |
14B从零训。32B从Qwen3-32B续训64B token，序列长度32K。
长上下文训练里注意力后半段是真金白银。14B在32K上少23.1%总训练FLOPs，不是只报kernel加速；128K反向3.0倍对step time更直接。Decode只有1.6倍，因为QK发现仍二次，推理时后半段占比没训练反向那么大。
MoBA、DSA、NSA先选support再算注意力，MALA不在打分前丢掉交互，associative recall几乎贴着FullAttn，固定预算方法掉了大半。代价是QK省不掉。短上下文预训练只省2.5%，收益明显偏长序列。
代码开源名为flash-sparse-attention。默认τ=1，训练推理同一规则，不用给每层每头单独校阈值。适合已经在用FlashAttention风格fused kernel、上下文往32K到128K走、又不想换稀疏拓扑的团队。这是执行层的常数因子优化，不是新的注意力函数。
QK仍是O(n²)。论文写明support-sparse方法还可以对排除区域跳过QK，MALA做不到。4K预训练只省2.5%，短序列别指望靠它改scaling曲线。
反向support嵌在前向里，但不是前向的精确反传；漏掉的梯度贡献还依赖dO和V，保真靠测量不是靠证明。前向漏质量有界τ/(1+τ)，τ=1时理论上限50%，实测千分之几，这个界很松。tile粒度保守：query block里一行不过关，整块都留。
KV cache原样保留，不减推理显存，只减算子工作集。Decode加速1.6倍，serving的大头仍可能是cache和带宽。32B是Qwen3-32B续训，不是从零用MALA训到32B。Associative recall没插distractor，测的是绑住的256对还在不在，不是噪声里找针。算子benchmark不含QKV投影、MLP、optimizer和流水线通信，端到端step加速会小于2.2倍和3.0倍。