Skip to content
Preprint

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

Jul 2026 · 0 citations · 41 references
Computer Science

TL;DR

CoSA is proposed, 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) and achieves a 4.93% attention speedup and reduces end-to-end Time-to-First-Token by 2.53% with negligible performance degradation.

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. Code is available at https://github.com/Tencent/AngelSlim.

View source

Similar papers

Preprint Jul 2026

RIS-Kernel: A Model-Agnostic Architecture for Long-Context LLM Inference via Sparse Attention

In controlled evaluations at 32,768 tokens, RIS-Stochastic at 1% density and 70 ensemble seeds achieves 75.00% accuracy, outperforming the native dense baseline, demonstrating that sparse attention acts as a regularizer: low density over multiple seeds filters out sequence-level noise, whereas higher density reintroduces distractor noise.

A. R. Santos · 1 citation
Preprint Jul 2026

LiteTopK: Exploiting the Curse of Dimensionality for a Fused Indexer-TopK Kernel in Long-Context Sparse Attention

Indexer-TopK, the operation to compute the scores and select the top-k candidates, is widely used by sparse attention algorithms in large language models and vector retrieval in recommendation systems and vector databases. However, existing GPU-based Indexer-TopK kernels like DeepSeek Sparse Attention (DSA) remain inefficient due to excessive global memory traffic, costly synchronization, and prohibitive memory overhead. In this study, inspired by the curse of dimensionality phenomenon, we first observe that sparse attention scores exhibit a score concentration phenomenon, where scores tend to fall within a narrow range. Based on this observation, we propose LITETOPK, an efficient fused Indexer-TopK kernel. LITETOPK first samples a small subset of data to estimate query-data score ranges, then partitions candidates into bins accordingly. This organization allows the LITETOPK kernel to maintain a tight approximate threshold online, write back only promising candidates, reduce unnecessary I/O and memory overhead while preserving exact Top-k correctness. Building on LITETOPK, we further propose LITEDSA, which exploits the similarity of top-k candidate sets among neighboring tokens. LITEDSA packs neighboring tokens'candidates for joint computation and masks out extra scores for each query, thereby reducing memory traffic while preserving correctness. Experimental results in a real-world deployment environ ment with eight B200 GPUs show that LITETOPK+LITEDSA accelerates the prefill stage of GLM 5.2 by 1.35x, with no performance loss and lower memory overhead.

Ziqi Yin, Jianyang Gao, Peiqi Yin et al. · 0 citations
#small language model Preprint Aug 2026

BF1: A Causal Dyadic Sparse-Attention Retrofit for Efficient Long-Context Transformers

The results establish BF1 as a reproducible sparse operator and selective retrofit primitive with real long-context systems value, and evaluates numerical correctness, selected-interaction scaling, kernel performance, partial-model inference, and matched next-token language modeling.

Hina Dixit · 0 citations
Preprint Jul 2026

COBS: Cumulant Order Block Sparse Attention

COBS (Cumulant Order Block Sparse Attention), an attention method that builds on NSA, incorporating a novel selector that stores a compressed second-order statistic per block, which achieves lower position-wise negative log-likelihood (NLL) than dense attention in the authors' comparison.

Alexander Tian, Aditya Ghai, S. Neelam et al. · 1 citation · ⚡1
Preprint Jul 2026

Prox: Training-Free FFN Activation Sparsity via Approximate Intermediate-Channel Salience in LLMs

Feed-forward networks (FFNs) dominate memory traffic and computation in large language model (LLM) inference, making them a primary target for activation sparsification. However, existing training-free methods suffer substantial model-quality degradation at high sparsity due to limitations in their channel-selection strategies. We observe that the SwiGLU intermediate state provides a highly effective channel-selection signal, but obtaining it requires costly dense computation. To address this, we present \emph{Prox}, a two-stage training-free framework for sparse SwiGLU FFNs. Prox hinges on the key insight: sparse execution requires only the channel mask induced by the intermediate state, which can be constructed from the magnitude ranking of its entries rather than their exact values. Specifically, Stage 1 uses input sparsity and quantized proxy weights to construct a shared mask; Stage 2 computes the selected channels exactly, enabling sparse execution of all three projections. Across ten LLMs from six model families, Prox outperforms training-free baselines at all sparsity levels, achieves up to a $1.99\times$ end-to-end decoding speedup at 70\% FFN sparsity, and is compatible with quantization and sparse attention.

Jinyi Liu, Wei Chen, Pengyu Chen et al. · 0 citations
Preprint Aug 2026

SPADE: An Input-Adaptive Sparse Attention Engine for Fast Video Diffusion Models Inference

SPADE is presented, a training-free sparse-attention engine of three parts: a specification defining 3D blocking candidates and formalizing dynamic masks via Summarizer/Estimator expressions, and an executor with low-overhead index search, flash block-sparse attention, and kernel grouping.

Shanghao Liu, Renze Chen, Size Zheng et al. · 0 citations