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

arXiv: https://arxiv.org/abs/2607.25291 · HTML: https://mlsanity.com/paper/2607.25291/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 41 references that are in the map.
- First posted on arXiv 28 Jul 2026.

## Builds on (most cited first, 30 of 35)

- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [GLM-5: from Vibe Coding to Agentic Engineering](https://mlsanity.com/paper/2602.15763/) (2026)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference](https://mlsanity.com/paper/2406.10774/) (2024)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://mlsanity.com/paper/2502.11089/) (2025)
- [MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention](https://mlsanity.com/paper/2407.02490/) (2024)
- [DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads](https://mlsanity.com/paper/2410.10819/) (2024)
- [FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving](https://mlsanity.com/paper/2501.01005/) (2025)
- [MoBA: Mixture of Block Attention for Long-Context LLMs](https://mlsanity.com/paper/2502.13189/) (2025)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [XAttention: Block Sparse Attention with Antidiagonal Scoring](https://mlsanity.com/paper/2503.16428/) (2025)
- [FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference](https://mlsanity.com/paper/2502.20766/) (2025)
- [SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference](https://mlsanity.com/paper/2502.18137/) (2025)
- [SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization](https://mlsanity.com/paper/2411.10958/) (2024)
- [FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling](https://mlsanity.com/paper/2603.05451/) (2026)
- [TokenSelect: Efficient Long-Context Inference and Length Extrapolation for LLMs via Dynamic Token-Level KV Cache Selection](https://mlsanity.com/paper/2411.02886/) (2024)
- [Retrieval Augmented Generation and Understanding in Vision: A Survey and New Outlook](https://mlsanity.com/paper/2503.18016/) (2025)
- [A-RAG: Scaling Agentic Retrieval-Augmented Generation via Hierarchical Retrieval Interfaces](https://mlsanity.com/paper/2602.03442/) (2026)
- [SageAttention2++: A More Efficient Implementation of SageAttention2](https://mlsanity.com/paper/2505.21136/) (2025)
- [HySparse: A Hybrid Sparse Attention Architecture with Oracle Token Selection and KV Cache Sharing](https://mlsanity.com/paper/2602.03560/) (2026)
- [ProxyAttn: Guided Sparse Attention via Representative Heads](https://mlsanity.com/paper/2509.24745/) (2025)
- [FlashPrefill: Instantaneous Pattern Discovery and Thresholding for Ultra-Fast Long-Context Prefilling](https://mlsanity.com/paper/2603.06199/) (2026)
- [SSA: Sparse Sparse Attention by Aligning Full and Sparse Attention Outputs in Feature Space](https://mlsanity.com/paper/2511.20102/) (2025)
- [Elastic Attention: Test-time Adaptive Sparsity Ratios for Efficient Transformers](https://mlsanity.com/paper/2601.17367/) (2026)
- [Stem: Rethinking Causal Information Flow in Sparse Attention](https://mlsanity.com/paper/2603.06274/) (2026)
- [Switch Attention: Towards Dynamic and Fine-grained Hybrid Transformers](https://mlsanity.com/paper/2603.26380/) (2026)

## Built on by (most cited first, 1 of 1)

- [FoldAttention: Declared-Reference Softmax for Fast Decode and Deterministic Backward](https://mlsanity.com/paper/2609.33410/) (2026)

## Abstract (arXiv)

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× attention speedup and reduces end-to-end Time-to-First-Token by 2.53× under a context length of 128K with negligible performance degradation. Code is available at https://github.com/Tencent/AngelSlim.