ATFlash：基于 RoPE 波长的注意力窗口，实现高效 LLM 推理

arXiv 发布论文 ATFlash，提出基于 RoPE 波长分段的注意力窗口，按频率对查询-键内积剪枝。在 Qwen2.5-0.5B 和 Llama-3.2-3B 上，原生上下文长度内剪除 37%-48% 内积项，top-1 匹配率保持 96%-98%，KL 散度在 10^-3 nat 量级。该窗口可作为 FlashAttention-4 和 FlashInfer 的补丁，在 RTX PRO 6000 上 Llama 模型加速最高 1.29 倍，且加速比随上下文长度增长。

此前动态稀疏注意力方法（如 MInference）的剪枝率依赖输入，且实现复杂。ATFlash 提出与输入无关、闭合形式（对数规模）的剪枝，简化预测和部署。同时，它作为一个正交操作，可叠加在其他稀疏方法之上，提升长上下文推理的计算和内存效率，在模型性能和加速间取得平衡。

RoPE 注意力分数可分解为多个 2D 旋转频率对之和，每个频率对的波长决定了其区分位置的最大距离。因此，对超过波长比例距离的查询-键内积进行剪枝，既保留了低频对的远程信息，又减少了高频对的冗余计算。该剪枝在频率层面进行，与 token 级动态剪枝相互独立，可组合使用。

该技术可降低长上下文 LLM 推理的计算和内存开销，加速比随上下文长度增长（128K 时 1.29 倍）。作为 FlashAttention-4 和 FlashInfer 的补丁，对现有推理库改动小，便于集成。产品侧预计可降低长文本服务的延迟和成本，提升用户体验。

• 在长文本处理产品中集成 ATFlash 补丁，以提升推理吞吐并降低延迟。
• 与现有动态稀疏方法结合，探索进一步提升剪枝率和速度。
• 针对超长上下文（如 1M）场景，利用 ATFlash 的输入无关剪枝特性简化资源规划。
• 基于 ATFlash 的封闭形式剪枝率，开发可预测性能的成本模型。

• 论文中 Qwen2.5-7B-1M 端到端剪枝 57%，但具体端到端速度提升未在摘要中说明。
• 在更长上下文和更大模型上的效果是否保持？
• ATFlash 在解码阶段（FlashInfer）的加速效果如何？
• 与其他剪枝方法组合后的实际收益是否存在上限？