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

As context lengths in Large Language Models (LLMs) grow, the quadratic computational cost of self-attention becomes a primary bottleneck for inference. Existing sparse attention methods typically decouple the proxy (which selects blocks to compute) from the kernel (which executes the computation). This separation often leads to suboptimal performance: proxies become unreliable at tight budgets, and kernels remain constrained by rigid, non-adaptive execution orders. CoSA addresses this by proposing a proxy-kernel co-design that treats the proxy and kernel as a unified system.

CoSA introduces a two-stage, training-free sparse attention mechanism:

1. Kernel-Aware Proxy (KAP): Instead of emitting a simple binary mask, KAP generates an ordered mask. It identifies 'Have Rowmax' (HRM) blocks—those containing the most salient attention scores—and prioritizes them in the execution sequence. This ensures that the most critical information is processed early.

2. Ordered-Skipping Kernel (OSK): The kernel consumes this ordered mask to perform two-stage sparsification. First, it uses the mask to physically jump over redundant KV pages, avoiding unnecessary memory access and synchronization. Second, it performs in-kernel skipping based on online-softmax statistics. Because the kernel processes the most important blocks first (as dictated by KAP), the running maximum of attention scores updates rapidly, allowing the kernel to safely and aggressively skip remaining blocks that are deemed negligible.

CoSA demonstrates that aligning the proxy's selection logic with the kernel's execution order significantly improves the efficiency-fidelity trade-off. By reordering block traversal to prioritize HRM blocks, the kernel can perform more aggressive skipping without sacrificing accuracy. Experimental results across mainstream LLM backbones (e.g., Llama-3.1-8B, Qwen3-8B) show that CoSA achieves a 4.93x attention speedup and a 2.53x reduction in Time-to-First-Token at 128K context lengths, outperforming existing sparse attention baselines while maintaining high accuracy.