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

2026-07-28

AI Summary

This work addresses the quadratic complexity of self-attention, which hinders efficient long-context reasoning, and the tendency of existing sparse attention methods to inadvertently discard critical information under low computational budgets, leading to significant accuracy degradation. To overcome these limitations, the authors propose CoSA, a training-free two-stage sparse attention method that introduces the first co-design of proxy and attention kernel. CoSA employs a Kernel-Aware Proxy (KAP) to generate ordered masks and integrates an Ordered Skipping Kernel (OSK) that dynamically skips redundant attention blocks using online softmax statistics. Evaluated on 128K-token contexts, CoSA achieves a 4.93× speedup in attention computation and a 2.53× reduction in first-token latency, while incurring only negligible performance loss.

Abstract

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.

Problem

• long-context inference
• self-attention
• sparse attention
• computational cost
• model accuracy

Innovation

• Sparse Attention
• Proxy-Kernel Co-Design
• Long-Context Inference
• Training-Free Acceleration
• Ordered-Skipping Kernel