arXiv:2602.22575cs.LGcs.AI2026-02ACL

通过在线排列实现稀疏注意力的早期停止,大幅提升长文本推理效率。

S2O: Early Stopping for Sparse Attention via Online Permutation

论文配图:S2O: Early Stopping for Sparse Attention via Online Permutation
图 1 · 摘自论文原文
  • 基于在线排列加载高优先级块,动态跳过低贡献内容。
  • 在128K上下文下,前缀计算密度降低3.31倍,误差可控。
  • 适合需要长序列推理的大型语言模型部署场景。

注意力机制随序列长度呈二次增长,根本限制了长上下文推理。现有块粒度稀疏化虽能降低延迟,但粗粒度块导致稀疏性天花板,难以进一步优化。本文提出S2O,通过在线排列实现稀疏注意力的早期停止。受内存系统虚拟-物理地址映射启发,S2O重构FlashAttention执行流程,使推理可加载非连续令牌,而非原顺序的连续段。基于注意力热图的细粒度结构,将显式排列转化为在线、索引引导的离散加载策略;仅需极轻量预处理和索引重映射开销,即可聚焦高优先级块。在此基础上,S2O引入早期停止规则:按重要性从高到低计算,当当前块得分低于阈值时提前终止,跳过剩余低贡献块,从而在可控误差预算下提升有效稀疏性并减少计算。结果表明,S2O显著提高实际稀疏性上限。在Llama-3.1-8B、128K上下文下,单算子MSE降低3.82倍(同稀疏度),前缀计算密度降低3.31倍(同MSE);同时保持端到端精度,获得7.51倍注意力加速和3.81倍端到端加速。

原文摘要 · Abstract (English)

Attention scales quadratically with sequence length, fundamentally limiting long-context inference. Existing block-granularity sparsification can reduce latency, but coarse blocks impose an intrinsic sparsity ceiling, making further improvements difficult even with carefully engineered designs. We present S2O, which performs early stopping for sparse attention via online permutation. Inspired by virtual-to-physical address mapping in memory systems, S2O revisits and factorizes FlashAttention execution, enabling inference to load non-contiguous tokens rather than a contiguous span in the original order. Motivated by fine-grained structures in attention heatmaps, we transform explicit permutation into an online, index-guided, discrete loading policy; with extremely lightweight preprocessing and index-remapping overhead, it concentrates importance on a small set of high-priority blocks. Building on this importance-guided online permutation for loading, S2O further introduces an early-stopping rule: computation proceeds from high to low importance; once the current block score falls below a threshold, S2O terminates early and skips the remaining low-contribution blocks, thereby increasing effective sparsity and reducing computation under a controlled error budget. As a result, S2O substantially raises the practical sparsity ceiling. On Llama-3.1-8B under a 128K context, S2O reduces single-operator MSE by 3.82$\times$ at matched sparsity, and reduces prefill compute density by 3.31$\times$ at matched MSE; meanwhile, it preserves end-to-end accuracy and achieves 7.51$\times$ attention and 3.81$\times$ end-to-end speedups.

稀疏注意力长序列推理早期停止FlashAttention

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。