💬 NLP

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

CoSA is a training-free, two-stage sparse attention framework that couples a Kernel-Aware Proxy with an Ordered-Skipping Kernel to dynamically optimize block selection and skipping, achieving significant inference speedups and reduced latency for long-context LLMs while maintaining high accuracy under tight computational budgets.

Original authors: Yufei Xue, Lin Niu, Hong Liu, Siran Liu, Hanyong Shao, Wei Liu, Guanghua Yu, Jianchen Zhu, Jun Zhang

Published 2026-07-29

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Yufei Xue, Lin Niu, Hong Liu, Siran Liu, Hanyong Shao, Wei Liu, Guanghua Yu, Jianchen Zhu, Jun Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to read a massive, 128,000-page encyclopedia to answer a single question. In the world of Artificial Intelligence, these "encyclopedias" are called Large Language Models (LLMs), and the "reading" process is how they understand context. The problem is that the standard way these models read is like a librarian who insists on reading every single page of the book, from the very first word to the very last, before they can even start thinking about your question. As the book gets longer, this "read-everything" approach becomes painfully slow and expensive, like trying to run a marathon while carrying a backpack full of bricks.

To fix this, scientists have tried to teach the librarian to be a bit more selective, a technique called "sparse attention." Instead of reading every page, the model tries to guess which pages are important and skips the rest. Usually, this involves two steps: first, a quick "proxy" (a fast guesser) looks at the book and marks the important pages with a simple "Yes" or "No" list. Second, a "kernel" (the actual worker) follows that list and does the heavy lifting. It's a decent system, but it has a flaw: when the book gets huge and you need to be very strict about which pages you skip to save time, the quick guesser starts making mistakes, and the worker just follows the bad list blindly. The result is that the model gets faster but starts forgetting important details, like a librarian who skips the climax of the story because they thought it was boring.

This is where a new method called CoSA (Co-Designed Sparse Attention) comes in. The researchers behind CoSA realized that the "guesser" and the "worker" were working in isolation, which caused the system to break down under pressure. They decided to redesign the whole process so the two work together as a team. Instead of just handing the worker a simple "Yes/No" list, the new "guesser" (called a Kernel-Aware Proxy, or KAP) hands over a priority list. It doesn't just say "read this page"; it says, "Read this page first, then this one, then this one."

Here is the magic trick: The worker (called an Ordered-Skipping Kernel, or OSK) uses this priority list to rearrange the order in which it reads the pages. By reading the most critical pages first, the worker builds up a "running score" of what's important very early on. Because it knows the most important stuff right away, it can confidently skip even more pages later in the process without getting confused. It's like if you were reading a mystery novel and the detective told you, "Read the first three chapters to find the killer, then you can safely skip the next fifty chapters of gardening tips."

The paper shows that this teamwork approach is a game-changer. When tested on models reading contexts as long as 128,000 tokens (roughly the size of a long novel), CoSA managed to speed up the "attention" part of the process by 4.93 times and cut the time it takes to generate the first word by 2.53 times. Crucially, it did this without the model losing its ability to understand the story or solve complex reasoning problems. The researchers found that by letting the proxy and the kernel talk to each other and share a specific order of operations, they could be much more aggressive about skipping work while still keeping the model smart. It turns out that in the world of AI, knowing when to read a page is just as important as knowing which page to read.

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

1. Problem Statement

The quadratic computational cost of self-attention mechanisms renders long-context inference prohibitively expensive for modern Large Language Models (LLMs). While sparse attention offers a remedy by selectively computing important query-key interactions, existing approaches face a fundamental trade-off between efficiency and fidelity, particularly under tight sparsity budgets.

Current methods generally fall into two categories:

Proxy-Driven Approaches: These use a lightweight proxy to predict a binary sparse mask before the kernel executes. While effective at moderate budgets, these proxies inevitably miss salient blocks as the budget tightens, leading to significant accuracy degradation.
In-Kernel Skipping Approaches: These use exact Softmax statistics within the kernel to skip negligible blocks. While accurate, they suffer from two limitations:
- Dense Proxy Requirement: They must compute the full Query-Key (QK) interaction for every block to make a skip decision, failing to eliminate the dominant quadratic QK cost.
- Conservative Skipping: Due to the "bucket effect" (where a single outlier row vetoes a skip) and the use of a running maximum (which is a lower bound of the global maximum) rather than a global maximum, these methods apply overly conservative skip rules, leaving potential speedups untapped.

The core bottleneck identified is the disconnection between the proxy and the backend kernel. Proxies generate static binary masks without considering execution order, while kernels execute in fixed orders that hinder optimal sparsity.

2. Methodology: CoSA

The authors propose CoSA, a two-stage, training-free block-sparse attention mechanism built on the principle of proxy-kernel co-design. CoSA couples a Kernel-Aware Proxy (KAP) with an Ordered-Skipping Kernel (OSK) to jointly optimize block selection and execution order.

2.1 Core Design Principles

Order Invariance: Leveraging the property that Online-Softmax (OSM) accumulation is independent of the order in which key blocks are visited, CoSA allows for arbitrary reordering of block processing without altering the final output.
Computation-Order Mask: Instead of a binary mask, CoSA generates a mask that prescribes the specific order in which KV pages should be visited.

2.2 Stage 1: Kernel-Aware Proxy (KAP)

The KAP operates under a moderate budget to select blocks and determine their visiting order.

Block Ranking: It downsamples queries and keys, computes approximate scores, and applies MaxPool over keys followed by MaxNorm. This process preserves the "rowmax" (the maximum value in a row) within blocks and converts scores into relative ranking contributions.
HRM Flag Generation: It identifies "Have Rowmax" (HRM) blocks—blocks containing the maximum value for a specific query row.
Flag-Grouped Sorting: The KAP constructs a computation-order mask ( $M_{KAP}$ ) that prioritizes visiting HRM blocks first, followed by other blocks sorted by their estimated scores. This ordering ensures that the kernel encounters the global maximum early in the traversal.

2.3 Stage 2: Ordered-Skipping Kernel (OSK)

The OSK consumes the $M_{KAP}$ mask to execute sparse attention with three distinct patterns:

Mask-Skipped: Blocks not selected by the KAP are physically jumped over via KV page remapping. This eliminates the control-flow and synchronization costs associated with skipped blocks in standard sparse kernels.
In-Kernel Skipped: For the blocks retained by Stage 1, the OSK performs a second-stage skip using exact online-softmax statistics. Because the HRM blocks were visited first (due to KAP's ordering), the running maximum rises quickly to the global maximum. This allows the skip condition (based on the difference between local and running max) to be satisfied more aggressively and safely, avoiding the "bucket effect."
Computed: Blocks that pass both filters are fully computed.

3. Key Contributions

Revealing the Disconnect: The paper identifies that the isolation between proxy design and kernel execution is the fundamental barrier to achieving high sparsity without accuracy loss.
KAP Design: A novel proxy that emits a computation-order mask rather than a binary mask. It selects blocks under a moderate budget and prescribes an execution order that prioritizes blocks containing rowmax values.
OSK Design: An optimized kernel that physically jumps over pages based on the KAP mask and performs aggressive in-kernel skipping by leveraging the improved running maximum statistics derived from the reordered traversal.
Empirical Validation: Extensive experiments demonstrating that CoSA achieves higher accuracy at lower budgets compared to state-of-the-art methods (MInference, FlexPrefill, XAttention) across various LLM backbones (Qwen3-8B, Llama-3.1-8B-Instruct) and benchmarks.

4. Experimental Results

Experiments were conducted on RULER (synthetic tasks) and LongBench-v2 (real-world reasoning) with context lengths ranging from 4K to 128K.

Accuracy: CoSA consistently outperforms existing sparse methods. On LongBench-v2 with Qwen3-8B, CoSA achieved the highest average accuracy (36.51%) while operating at the lowest budget (15%). It maintained robust performance in both "w/o CoT" (direct) and "w/ CoT" (reasoning) modes.
Efficiency:
- Attention Speedup: CoSA achieves a 4.93× speedup in attention computation at a 128K context length.
- End-to-End Latency: It reduces the Time-to-First-Token (TTFT) by 2.53× at 128K context with negligible performance degradation compared to dense attention.
Ablation Studies: The paper demonstrates that the combination of KAP (ordered selection) and OSK (page remapping + in-kernel skip) is synergistic. Using KAP alone improves accuracy; adding in-kernel skipping increases sparsity; and adding page remapping restores and improves accuracy while further tightening the budget.

5. Significance and Claims

The paper claims that CoSA represents a paradigm shift from treating the proxy and kernel as isolated components to a co-designed system. By bridging them with a computation-order mask, CoSA resolves the efficiency-fidelity trade-off that plagues current sparse attention methods.

The authors assert that:

Aggressive sparsity is achievable without accuracy loss: By front-loading HRM blocks, the kernel can skip more blocks safely, pushing the sparsity budget lower than previously possible.
Physical execution matters: Simply having a good mask is insufficient; the kernel must physically reorder memory access (via page remapping) to realize the theoretical speedups and avoid redundant synchronization costs.
Scalability: CoSA is particularly effective for long contexts (128K+), where the quadratic cost of attention dominates, offering substantial real-world latency reductions for applications like Retrieval-Augmented Generation (RAG) and autonomous agents.

The paper concludes by noting that while CoSA currently focuses on the prefilling stage, future work will explore extending these co-design principles to the decoding stage, which presents distinct computational patterns.