MassAlloc Attention：让Attention分配自己的计算

FullAttn 通常将可忽略的归一化质量分配给大部分因果分数空间，但密集内核在形成每个 QK 块后执行完整的后分数路径。我们引入了 MALA，一种融合注意力原语，它保留对每个法律因果交互的分数访问，并使用归一化贡献来分配分数后计算。前向使用其不断发展的在线softmax标准化器，而后向重用最终标准化器以仅使用标准注意状态来导出嵌套保留支持。共同的容忍度控制着训练和推理，允许适应性地保留工作。 MALA 减少了低贡献的分数后计算。 8K 的匹配工作研究分离了分布自适应分配的好处：在完全匹配的总分后工作下，MALA 接近每个实例的参考质量预言，平均省略质量为 0.0188% 与 0.0182%。在从 1K 到 32K 词元的上下文长度中，相同的容差保持了相对于参考的低输出和梯度误差。在更广泛的受控联想回忆比较中，随着上下文的增长，MALA 密切跟踪 FullAttn，在 8K 下达到 89.67% 的准确率，而 FullAttn 的准确率为 89.97%。在具有张量并行性的 128K 词元的注意力算子基准测试中，相对于 FullAttn，MALA 将训练期间的前向和后向延迟减少了 2.2 倍和 3.0 倍，将推理期间的解码延迟减少了 1.6 倍。在从 0.6B 到 14B 参数的标度律训练中，MALA 密切跟踪 FullAttn 的困惑度，同时减少总训练 FLOP。通过单独的持续训练得到的 14B 模型和 32B 模型获得了与 FullAttn 相当的知识、推理和长上下文检索分数。这些结果表明，根据归一化注意力贡献分配后评分计算可以保留 FullAttn 的评估能力，同时减少注意力计算。