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

cs.AIOct 20, 2026

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

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

Organizations: Department of Computer Science, Rice University, TX, USA. · Department of Electrical Engineering and Computer Sciences, UC Berkeley, CA, USA. · Department of Computer Science, Stanford University, CA, USA. · Department of Computer Science and Engineering, Indian Institute of Technology Bombay, India.

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 $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.

Figures & tables

Panel of Figure 1 — Figure 1: Relative $L_{2}$ error when varying the top- $k$ :sampling budget allocation at $50\times$ sparsity (left) and varying sparsity with a fixed 3:2 top- $k$ :sampling budget split (right).

Panel of Figure 2 — Figure 2: Average attention error across different context lengths and sparsity levels for Qwen3.5-9B and Qwen3.5-4B on LOFT [ 22 ] .

Table 3

Sparsity	Batch	FlashInfer	Reference	SSA	vs. FlashInfer	vs. reference
		Latency (ms)			SSA speedup
50 $\times$	1	1.41	2.99	0.79	1.79 $\times$	3.80 $\times$
	8	11.90	7.18	4.27	2.78 $\times$	1.68 $\times$
	16	24.47	12.57	8.26	2.96 $\times$	1.52 $\times$
20 $\times$	1	1.41	3.34	1.01	1.40 $\times$	3.32 $\times$
	8	11.91	8.74	5.76	2.07 $\times$	1.52 $\times$

Table 2: Kernel optimizations make SSA faster than dense decoding even at batch size 1. We report decode-time attention-layer latency at 1M context length on a single H200 GPU. Reference implements the same algorithm without the kernel optimizations described in Section 4.3 . At batch size 1, it reaches only 0.37-0.47

\times

the speed of FlashInfer and surpasses FlashInfer only at batch size 8. In contrast, optimized SSA is faster than FlashInfer across all measured batch sizes and sparsity levels. FlashInfer latencies across the two runs agree within 1.2%.

	Qwen3.5-9B				Qwen3.5-4B
Sparsity	Oracle-top- $k$	PQ	vAttn(PQ)	SSA	Oracle-top- $k$	PQ	vAttn(PQ)	SSA
Dense	66.3				54.2
5 $\times$	63.8	63.5	65.4	61.0	49.8	49.4	54.4	52.7
10 $\times$	59.4	53.8	62.7	64.0	47.9	41.5	51.9	51.0
20 $\times$	55.4	49.4	58.3	63.5	42.7	33.8	50.2	53.1
50 $\times$	48.3	36.7	35.8	59.8	36.5	25.4	27.5	50.0

Table 3: AIME 2025 avg@16 accuracy (left). LongBench performance with Qwen3.5-27B (right).

Figure 3: SSA accelerates decode at long contexts, with larger gains at higher batch sizes. Decode attention speedup on a single H200 (16 query heads, 4 KV heads, head dimension 256, fp16 KV). (a) At batch size 1, SSA overtakes FlashInfer beyond 256K context and reaches 1.79 $\times$ speedup at 1M and 50 $\times$ sparsity. (b) At 1M context, speedup increases with batch size, reaching 2.96 $\times$ at batch 16 and 50 $\times$ sparsity. (c) Overhead of two-level residual sampling relative to 1,024-token chunk-only sampling as the number of draws $m$ increases.

Figure 4: Illustration of Stratified Sparse Attention algorithm.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Average Subspan-EM and error across different context lengths and sparsity levels on LOFT.

Model	Sparsity	SSA-Theory	SSA-Practical
Llama-3.1-8B	50 $\times$	81	83
Llama-3.1-8B	20 $\times$	88	88
Ministral-3-8B	50 $\times$	82	84
Ministral-3-8B	20 $\times$	89	88

Appendix

Table 4: Comparison of theoretical and practical SSA variants on average RULER-HARD-32K performance.

Model	Dense	PQCache	vAttn(PQ)	Quest	HashAttn	Double Sparsity	MagicPig
Qwen3-4B-Instruct-2507	88.67	85.80	85.67	42.03	–	39.29	18.72
Llama-3.2-3B-Instruct	66.25	39.30	39.59	34.11	–	31.21	17.22
Llama-3.1-8B-Instruct	87.89	68.95	64.42	64.17	49.48	44.62	16.28

Appendix

Table 5: Average performance on RULER-HARD-32K across different models at 50

\times

sparsity.

Dataset	# Samples	Subsets
RULER-HARD	100	fwe, qa1, qa2, niah_multikey_2, niah_multikey_3, vt
LongBench	200	NrtvQA, Qasper, MF-en, HotpotQA, 2WikiMQA, Musique, DuReader, GovRep, QMSum, MultiNews, VCSum, TREC, TriviaQA, SAMSum, PCount, PR-en, LCC, RBP
LOFT	100	HotpotQA, Musique, NQ, Qampari, Quest
AIME	16	–

Appendix

Table 6: Evaluation datasets and subsets used in our experiments.

Sparsity	Context	Method	FWE	VT	NM2	NM3	QA1	QA2	Avg.
50 $\times$	32K	Dense	93.33	99.2	100	100	88	52	88.76
		PQ	77.67	95.4	67	71	76	48	72.51
		vAttn (PQ)	83.67	94.6	52	58	66	46	66.71
		Oracle-top- $k$	81.00	97.8	87	76	79	52	78.80
		SSA	91.00	98.0	83	92	85	53	83.67
	64K	Dense	86.33	98.8	98	99	84	53	86.52

Appendix

Table 7: Performance on RULER-HARD with Llama-3.1-8B.

Sparsity	Context	Method	FWE	VT	NM2	NM3	QA1	QA2	Avg.
50 $\times$	32K	Dense	96.33	91.8	98	99	79	56	86.69
		PQ	92.67	60.2	91	56	73	52	70.81
		vAttn (PQ)	90.33	72.2	67	77	65	44	69.26
		Oracle-top- $k$	93.67	62.2	99	55	75	52	72.81
		SSA	95.67	88.6	92	90	72	54	82.05
	64K	Dense	74.33	99.8	93	96	75	55	82.19

Appendix

Table 8: Performance on RULER-HARD with Ministral-3-3B.

Sparsity	Method	NrtvQA	Qasper	MF-en	MF-zh	HotpotQA	2WikiMQA	Musique	DuReader	GovRep	QMSum
50 $\times$	Dense	20.49	29.64	38.94	42.97	34.83	34.64	26.69	25.92	31.91	21.96
	PQ	20.13	28.69	38.60	42.01	33.27	31.92	20.61	25.54	31.92	21.64
	vAttn(PQ)	21.97	25.33	39.38	43.53	32.00	30.44	21.45	26.91	30.68	22.42
	SSA	20.42	27.51	38.87	42.13	34.64	32.79	24.88	24.82	31.96	22.72
	Oracle-top- $k$	20.51	28.97	38.22	42.12	34.79	33.77	22.50	25.76	32.86	22.52
20 $\times$	PQ	20.61	29.55	38.41	41.76	34.51	33.38	23.90	25.74	31.63	22.31

Appendix

Table 9: Performance on LongBench with Qwen3.5-27B (10K-16K context).

Context	Method	FWE	VT	NM2	NM3	QA1	QA2	Avg.
32K	Dense	99.67	100	100	100	89	63	91.95
	PQ	97.33	100	100	100	88	62	91.22
	vAttn (PQ)	96.33	100	100	100	89	58	90.56
	Oracle-top- $k$	99.33	100	100	100	88	61	91.39
	SSA	99.67	100	100	100	89	60	91.45
64K	Dense	99.67	100	100	100	88	59	91.11

Appendix

Table 10: Performance on RULER-HARD with Qwen3.5-4B at 50

\times

sparsity.

Context	Method	FWE	VT	NM2	NM3	QA1	QA2	Avg.
32K	Dense	99.67	100	100	100	93	68	93.45
	PQ	94.00	100	100	100	91	67	92.00
	vAttn (PQ)	98.00	100	100	100	89	63	91.67
	Oracle-top- $k$	98.33	100	99	100	91	67	92.56
	SSA	100.00	100	100	100	93	67	93.33
64K	Dense	99.33	100	100	100	88	66	92.22

Appendix

Table 11: Performance on RULER-HARD with Qwen3.5-9B at 50

\times

sparsity.

Context	Sparsity	Method	FWE	VT	NM2	NM3	QA1	QA2	Avg.
32K	50 $\times$	Dense	94.00	100.00	99.00	100.00	80.00	63.00	89.33
		PQ	88.00	100.00	97.00	67.00	80.00	59.00	81.83
		vAttn (PQ)	89.33	98.80	85.00	87.00	64.00	48.00	78.69
		SSA	94.00	100.00	88.00	91.00	74.00	58.00	84.17
		Oracle-top- $k$	90.67	100.00	99.00	90.00	81.00	59.00	86.61
64K	50 $\times$	Dense	81.00	99.80	100.00	100.00	75.00	61.00	86.13

Appendix

Table 12: Performance on RULER-HARD with Ministral-3-8B.

Heavy Fraction $h$	20 $\times$	10 $\times$
0.0	40.2	42.5
0.2	61.0	63.8
0.4	63.5	64.0
0.6	62.7	62.7
0.8	63.8	63.1
1.0 (PQ)	49.4	53.8

Appendix

Table 13: Ablation on the heavy-hitter fraction

h

for SSA on AIME 2025. A fraction

h

of the sparse attention budget is allocated to deterministic heavy hitters, while the remaining

1-h

is allocated to importance sampling. Thus,

h=0

corresponds to pure importance sampling and

h=1

corresponds to PQCache (pure top-

k

retrieval).

Sparse Attention

Sparse attention reduces the token interactions computed by Transformers to lower the cost of long-context modeling and inference.

See all 243 papers

Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps
May 16, 2026Yanke Zhou, Yiduo Li, Hanlin Tang +6Self-Attention Long-Context Language Model Inference
Nanjing University · Alibaba Group
CSAttention: Centroid-Scoring Attention for Accelerating LLM Inference
Mar 30, 2026Chuxu Song, Zhencan Peng, Jiuqi Wei +1LLM Inference Acceleration Long-Context Language Model Inference
Department of Computer Science, Rutgers University · Oceanbase, Ant Group.
Near-Oracle KV Selection via Pre-hoc Sparsity for Long-Context Inference
Feb 9, 2026Yifei Gao, Lei Wang, Rong-Cheng Tu +3Self-Attention KV Caching
University College London, London, UK · Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences (CAS), Shenzhen 518055, China · The Chinese University of Hong Kong +1

May 16, 2026J/K move · Enter open · S save

Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps

Yanke Zhou, Yiduo Li, Hanlin Tang, Maohua Li+5

Nanjing University · Alibaba Group

Self-Attention Long-Context Language Model Inference LLM Inference Acceleration Sparse Attention Multi-Head Attention

Open paper AskPDF arXiv

Long-context inference in large language models is bottlenecked by the quadratic cost of full attention. Existing efficient alternatives often rely either on native sparse training or on heuristic token eviction, creating an undesirable trade-off among efficiency, training cost, and accuracy. In this work, we show that full-attention LLMs are already intrinsically sparse and can be transformed into highly sparse models with only minimal adaptation. Our approach is built on three observations: (1) only a small subset of attention heads truly requires full long-context processing; (2) long-range retrieval is governed primarily by a low-dimensional subspace, allowing relevant tokens to be retrieved efficiently with a 16-dimensional indexer; and (3) the useful token budget is strongly query-dependent, making dynamic top-

p

selection more suitable than fixed top-

k

sparsification. Based on these insights, we propose RTPurbo, which retains the full KV cache only for retrieval heads and introduces a lightweight token indexer for sparse attention. By exploiting the model's intrinsic sparsity, RTPurbo achieves sparsification with only a few hundred training steps. Experiments on long-context benchmarks and reasoning tasks show that RTPurbo preserves near-lossless accuracy while delivering substantial efficiency gains, including up to a 9.36

\times

prefill speedup at 1M context and about a 2.01

\times

decode speedup. These results suggest that strong sparse inference can be obtained from standard full-attention training without expensive native sparse pretraining.

Abstract

Sparse Attention

Explore similar work

Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps

CSAttention: Centroid-Scoring Attention for Accelerating LLM Inference

Near-Oracle KV Selection via Pre-hoc Sparsity for Long-Context Inference