arXiv:2607.25291cs.CL2026-07

通过代理与内核协同设计,显著加速长文本推理且保持高精度。

CoSA: Accelerating Long-Context Inference via Proxy-Kernel Co-Designed Sparse Attention

论文配图:CoSA: Accelerating Long-Context Inference via Proxy-Kernel Co-Designed Sparse Attention
图 1 · 摘自论文原文
  • 代理预测访问顺序,内核按序跳过冗余块,实现动态稀疏计算。
  • 128K上下文下速度提升4.93倍,首字延迟降低2.53倍,性能几乎无损。
  • 无需训练,适配主流大模型,适合资源受限的长文本生成场景。

自注意力的二次复杂度使长上下文推理成本高昂,基于代理的块稀疏注意力成为实用解决方案。现有方法通常依赖代理预测二值稀疏掩码,再由内核根据掩码执行稀疏计算。该方式在中等预算下有效,但预算收紧时,代理易遗漏关键块,而内核仅机械应用掩码,导致模型精度明显下降。本文提出CoSA,一种无需训练的代理-内核协同设计稀疏注意力方法,结合核感知代理(KAP)与有序跳过内核(OSK)。第一阶段,KAP在中等预算下选择关键块,并生成规定内核内部循环访问顺序的有序掩码;第二阶段,OSK基于在线Softmax统计,在更紧预算下跳过更多块。在主流LLM骨干网络和长上下文基准上,CoSA在更低预算下实现更高精度。令人印象深刻的是,当上下文长度为128K时,CoSA实现4.93倍注意力加速,端到端首字延迟降低2.53倍,性能损失可忽略。代码已开源:https://github.com/Tencent/AngelSlim。

原文摘要 · Abstract (English)

The quadratic cost of self-attention makes long-context inference prohibitively expensive, and proxy-based block-sparse attention has become a practical remedy. Existing methods typically rely on a proxy to predict a binary sparse mask and a kernel to consume this mask and perform sparse attention computation. Such an approach is effective under moderate budgets. However, as the budget tightens, the estimated proxy inevitably drops some salient blocks, while the kernel can only apply the sparse mask mechanically, leading to an evident drop in model accuracy. We propose CoSA, a two-stage training-free Sparse Attention under proxy-kernel CO-design, which couples a Kernel-Aware Proxy (KAP) with an Ordered-Skipping Kernel (OSK). In the first stage, the KAP selects blocks under a moderate budget and produces an ordered mask that prescribes the order in which KV pages are visited in the kernel inner loop. In the second stage, the OSK applies this mask and skips more blocks under a tightened budget given online-softmax statistics. Across mainstream LLM backbones and long-context benchmarks, CoSA attains higher accuracy at lower budgets. Impressively, CoSA achieves a 4.93$\times$ attention speedup and reduces end-to-end Time-to-First-Token by 2.53$\times$ under a context length of 128K with negligible performance degradation. Code is available at https://github.com/Tencent/AngelSlim.

稀疏注意力长文本生成推理加速大模型优化

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