Newest papers

Stratified Sparse Attention: Reducing the Sparse-to-Dense Gap at Extreme Sparsity

Sahil Joshi, Prithvi Dixit, Aryan Katiyar, Agniva Chowdhury, Joseph E. Gonzalez, Anshumali Shrivastava, Aditya Desai, Kumar Krishna Agrawal · 2026 · arXiv (Cornell University)

DOI: 10.48550/arxiv.2610.28441 · cited by 0 works (OpenAlex) · open-access PDF

Chinese

大型代码转换器框架的优化在较长的背景下越来越容易受到长长的关键值缓存(KV)缓存的成本的影响。 Sparse 注意减少了通过访问仅有一小部分缓存的成本,但现有的方法仍然面临着确定性检索之间的差距,捕捉高分数的代币,但排除剩余关注量和采样,从而提高了更高的差异的覆盖率。混合式方法将顶级 $k 采样与统一的采样结合起来,但仍然显示出从密集的关注到侵略性差异的显著差距。我们引入了 \emph{Stratified Spartense Attention (SSA) },一个统一的框架,将确定性检索与分数引导的采样结合起来。SSA 决定性地注意到高分

Machine translation (Meta M2M100); the English original follows.

Abstract

Large Language Model inference over long contexts is increasingly bottlenecked by the cost of attending to a growing Key-Value (KV) cache. Sparse attention reduces this cost by accessing only a fraction of the cache, but existing methods face a trade-off between deterministic retrieval, which captures high-scoring tokens but discards residual attention mass, and sampling, which improves coverage at the cost of higher variance. Hybrid methods combine top-$k$ retrieval with uniform sampling, yet still exhibit a substantial gap from dense attention at aggressive sparsity. We introduce \emph{Stratified Sparse Attention (SSA)}, a unified framework combining deterministic retrieval with score-guided sampling. SSA deterministically attends to high-scoring tokens while sampling the residual using proxy attention scores and correcting for their inclusion probabilities, thereby recovering attention mass beyond the retrieval set. SSA encompasses top-$k$, top-$p$, and hybrid top-$k$+sampling as special cases. Across RULER, LongBench, LOFT, and AIME, and five models ranging from 3B to 27B parameters, SSA consistently improves over the baselines at matched sparsity. On RULER-HARD with Ministral-3-8B, SSA improves accuracy over the strongest practical baseline from 81.8\% to 84.2\% at 32K and from 63.8\% to 69.3\% at 128K under $50\times$ attention sparsity, recovering 94\% and 85\% of dense-attention performance, respectively. With optimized CUDA kernels, SSA achieves up to $2.96\times$ decode speedup over FlashInfer at 1M context, demonstrating that its accuracy gains translate into practical long-context inference acceleration. We open-source our code at https://github.com/skylight-org/SSA-Paper-Code.

Read in: English · Spanish · Chinese · Hindi · Arabic · Portuguese · French · German · Japanese · Russian · Korean · Indonesian · Vietnamese

Topics: Computer science, Inference, Residual, Ranging, Speedup, Proxy (statistics)

Sources: OpenAlex (CC0), arXiv, Crossref. Agent Exchange