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

German

Große Code-Kontext-Transformer-Modell-Inferenz in langen Kontexten wird zunehmend durch die Kosten der Teilnahme an einem wachsenden Key-Value (KV) -Cache blockiert. Sparse-Aufmerksamkeit reduziert diese Kosten, indem sie nur einen Bruchteil des Cache erreicht, aber bestehende Methoden sehen immer noch einen wesentlichen Unterschied zwischen deterministischem Abrufen, das hochrangige Token erfasst, aber übrigbleibende Aufmerksamkeitsmasse (SSA) verdrängt. Sparse-Aufmerksamkeit verbessert die Abdeckung auf Kosten höherer Varianz. Hybrid-Methoden kombinieren Top-$k$-Retry mit einheitlichem Sampling, aber die bestehenden Methoden zeigen immer noch eine erhebliche Lü

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