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

title

authors

Sahil Joshi, Prithvi Dixit, Aryan Katiyar, Agniva Chowdhury, Joseph E. Gonzalez, Anshumali Shrivastava, Aditya Desai, Kumar Krishna Agrawal

date

2026-10-20

venue

arXiv (preprint, not peer reviewed)

arXiv:2610.28441

affiliation

Skylight

upvotes

7 (Hugging Face Daily Papers, 2026-11-01)

links

abs pdf hf code

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 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 50times attention sparsity, recovering 94\% and 85\% of dense-attention performance, respectively. With optimized CUDA kernels, SSA achieves up to 2.96times 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.

@misc{joshi2026stratified,
  title={Stratified Sparse Attention: Reducing the Sparse-to-Dense Gap at Extreme Sparsity},
  author={Sahil Joshi and Prithvi Dixit and Aryan Katiyar and Agniva Chowdhury and Joseph E. Gonzalez and Anshumali Shrivastava and Aditya Desai and Kumar Krishna Agrawal},
  year={2026},
  eprint={2610.28441},
  archivePrefix={arXiv},
  url={https://arxiv.org/abs/2610.28441}
}