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

May 16, 2026

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

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.

Yanke Zhou, Yiduo Li, Hanlin Tang +6

Nanjing University · Alibaba Group

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

Ask PDF

Mar 30, 2026

CSAttention: Centroid-Scoring Attention for Accelerating LLM Inference

Long-context LLMs increasingly rely on extended, reusable prefill prompts for agents and domain Q&A, pushing attention and KV-cache to become the dominant decode-time bottlenecks. While sparse attention reduces computation and transfer costs, it often struggles to maintain accuracy at high sparsity levels due to the inherent distribution shift between Queries and Keys. We propose Centroid-Scoring Attention (CSAttention), a training-free sparse attention method optimized for high-throughput serving of reusable contexts. CSAttention adopts a storage-for-computation strategy tailored to the offline-prefill/online-decode setting: it front-loads computation into a one-time offline prefill phase that can be amortized across multiple queries, while aggressively optimizing per-step decoding latency. Specifically, CSAttention constructs query-centric lookup tables during offline prefill, whose size remains fixed during decoding, and enables online decoding to replace full-context scans with efficient table lookups and GPU-friendly score accumulation. Extensive experiments demonstrate that CSAttention achieves near-identical accuracy to full attention. Under high sparsity (95%) and long-context settings (32K-128K), CSAttention consistently outperforms state-of-the-art sparse attention methods in both model accuracy and inference speed, achieving up to 4.6x inference speedup over the most accurate baseline at a context length of 128K.

Chuxu Song, Zhencan Peng, Jiuqi Wei +1

Department of Computer Science, Rutgers University · Oceanbase, Ant Group.

LLM Inference Acceleration Long-Context Language Model Inference Attention Mechanisms Sparse Attention

Ask PDF

Feb 9, 2026

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

A core bottleneck in large language model (LLM) inference is the cost of attending over the ever-growing key-value (KV) cache. Although near-oracle top-k KV selection can preserve the quality of dense attention while sharply reducing computation and bandwidth, existing sparse methods generally rely on posterior heuristics, i.e., selectors conditioned on observed attention or proxy scores. Such conditioning introduces posterior bias: it tends to distort true token importance and miss salient tokens, thereby impairing long-range reasoning. To tackle this problem, we propose Pre-hoc Sparsity (PrHS), which selects KV entries before attention scoring and provides explicit accuracy control. Let the attention mass of discarded entries be delta (the dropped mass). Through a marginal-to-mutual-information analysis, we derive an upper bound on the mutual-information loss that depends only on the dropped mass. This relation explains failure modes of posterior heuristics and enables verifiable guarantees by controlling the dropped mass in advance. Within PrHS, we instantiate three orthogonal pre-hoc selectors along the axes of time, depth, and layer. Extensive experiments on LLaMA and Mistral families validate PrHS. Across GSM8K and CoQA, PrHS reduces retrieval overhead by over 90%, achieving 3x higher retrieval sparsity than HShare at matched or better accuracy. It incurs under 1% average degradation on LongBench, lowers attention FLOPs by about 15% versus prior sparse baselines, and yields a 9.9x speedup in attention-operator latency and 2.8x higher throughput on NVIDIA A100-80GB GPUs than the dense baseline.

Yifei Gao, Lei Wang, Rong-Cheng Tu +3

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

Self-Attention KV Caching Long-Context Language Model Inference KV-Cache Management LLM Inference Acceleration Sparse Attention+1

Ask PDF

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