SeerAttention-R: 用于长程推理的稀疏注意力自适应

我们引入了 SeerAttention-R，一个专门为推理模型的长解码定制的稀疏注意力框架。 SeerAttention-R 扩展自 SeerAttention，保留了通过自蒸馏门控机制学习注意力稀疏性的设计， 同时去除了查询池化以适应自回归解码。SeerAttention-R 具有轻量级即插即用门控， 灵活且可以轻松集成到现有预训练模型中，而无需修改原始参数。 我们证明，SeerAttention-R 仅在 0.4B token 上进行训练，在 AIME 基准测试中， 在大稀疏注意力块大小（64/128）下，使用 4K token 预算仍能保持接近无损的推理准确性。 使用 TileLang，我们开发了一个高度优化的稀疏解码内核， 在 H100 GPU 上，在 90% 稀疏度下，实现了相对于 FlashAttention-3 高达 9 倍的接近理论值的加速。 代码可在以下网址获取：https://github.com/microsoft/SeerAttention。