Source-linked AI summary
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
TL;DR
Long-context inference is increasingly constrained by attention over a growing KV cache, while sparse methods trade retrieval of high-scoring tokens against broader but higher-variance sampling. SSA combines deterministic selection with score-guided sampling of residual tokens and corrects for their inclusion probabilities. Across the reported evaluations, it consistently improves over sparse-attention baselines and narrows the gap to dense attention.
Problem
Long-context decoding repeatedly attends to a growing KV cache, while existing sparse methods trade high-scoring-token retrieval against broader sampling and still leave a dense-attention gap at aggressive sparsity.
Source: §1 · pp. 1–2Method
SSA retains high-proxy-score tokens and samples residual tokens according to proxy scores, correcting selected-token contributions by their inverse inclusion probabilities.
Source: §4 · p. 5, §4.2 · p. 6Results
SSA consistently improves over deterministic, sampling-based, and hybrid sparse-attention baselines across RULER, LongBench, LOFT, and AIME evaluations.
Source: §1 · p. 3Takeaways & Limitations
SSA encompasses top-k, top-p, and hybrid top-k+sampling and is compatible with approximate top-k methods that provide proxy scores.
Source: §1 · p. 3, §6 · p. 11The optimized GPU implementation uses a fixed heavy fraction; an equally efficient implementation of the adaptive allocation remains future work.
Source: §4.3 · p. 7
Abstract
from arXiv · showhide
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.
1 Introduction
Long-context decoding becomes increasingly memory-bandwidth bound as each generated token attends to a growing KV cache. SSA combines deterministic retrieval with score-guided sampling to improve sparse attention’s coverage, and the paper reports consistent gains over sparse-attention baselines across its evaluations.
1 Introduction
- Long-context decoding becomes memory-bandwidth bound as each step transfers the growing KV cache from HBM to on-chip SRAM.Sparse attention reduces memory traffic and latency by accessing only a fraction of the cache.Source: §1 · pp. 1–2
- Score-guided importance sampling achieves lower relative L2 error than uniform sampling, according to Figure 1.Figure 1 varies the top-k:sampling budget allocation at 50× sparsity and varies sparsity with a fixed 3:2 split.Source: Figure 1 · p. 2, §1 · p. 2
- Deterministic retrieval can miss residual attention mass, while sampling broadens coverage but introduces variance; hybrid methods still leave a gap to dense attention at aggressive sparsity.SSA retains high-scoring tokens and samples the residual according to proxy attention scores.Source: §1 · p. 2
- SSA consistently improves over deterministic, sampling-based, and hybrid sparse-attention baselines across four benchmarks and five models.The framework also provides theoretical bounds and CUDA/Triton kernels for end-to-end decoding speedups.Source: §1 · p. 3
- SSA unifies deterministic and stochastic token selection, with top-k, top-p, and hybrid top-k+sampling as special cases.The paper also presents theoretical bounds and efficient CUDA/Triton kernel support.Source: §1 · p. 3
2 Related Work
Related work addresses sparse attention through KV-cache compression, approximate retrieval, and sampling or hybrid selection. SSA’s distinction from uniform hybrid sampling is to guide residual-token selection with approximate attention scores.
Head-level compression and KV eviction
Approximate top-k/top-p methods
- Approximate top-k methods recover high-attention tokens but discard all tokens outside the retrieval set, potentially losing residual-tail mass at aggressive sparsity.Tactic [50] extends retrieval to top-p by estimating how many high-scoring tokens meet a target attention-mass threshold.Source: §2 · p. 4
Sampling-based and Hybrid methods
3 Background and Motivation
Sparse attention approximates dense scaled dot-product attention by selecting fewer KV-cache tokens. Deterministic top-k selection retains high-scoring tokens, while importance sampling uses probability-weighted samples to estimate the dense output.
3.1 Sparse Attention
- Dense attention computes a query’s scaled dot-product attention over the KV cache, whose entries have scores s_i and weights proportional to exp(s_i).The output is the weighted combination of the value vectors.Source: §3.1 · p. 4
- Deterministic top-k selection retains a heavy set H of cached tokens, but tokens outside H do not contribute to the sparse output.Source: §3.1 · p. 4
3.2 Importance Sampling
- Importance-based sparse attention weights sampled tokens by inverse inclusion probability to estimate the dense output; importance sampling can lower variance when its proposal favors high-contribution samples.The ratio p(x)/q(x) corrects for drawing from a proposal distribution rather than the target distribution.Source: §3.1 · pp. 4–5, §3.2 · p. 5
4 Introducing Stratified Sparse Attention (SSA)
SSA combines deterministic retention of high-scoring tokens with score-guided sampling of the residual. Its algorithm allocates a fixed budget between these strata, while the analysis studies approximation error and the kernel section describes an implementation that currently uses a fixed heavy fraction.
4 Introducing Stratified Sparse Attention (SSA)
- SSA retains high-confidence heavy hitters and samples the residual according to proxy attention scores, balancing deterministic exploitation with stochastic exploration.This score-aware stratification is intended to preserve heavy hitters while providing stochastic coverage toward more promising residual tokens.Source: §4 · p. 5
4.1 Final Algorithm
- The algorithm sets a threshold on proxy probabilities to form the heavy set, then samples the residual and corrects each selected token’s exact score by its inclusion probability.The proxy distribution determines the split without a separate heavy-set-size choice; duplicate draws are retained only once, so the selected set has size at most k.Source: §4.1 · p. 5
4.2 Theoretical Insights
- The analysis treats k < N as the sampling case and identifies the error trade-off: larger residual mass and proxy mismatch increase error, while more residual draws reduce it at the Monte Carlo rate m^-1/2.When k = N, all tokens are retained and SSA returns dense attention exactly; the residual-mass and proxy-coverage quantities characterize the sampling error.Source: §4.2 · pp. 6–7
4.3 Kernel Design
- The optimized GPU implementation uses a fixed heavy fraction rather than the algorithm’s adaptive split, and an equally efficient adaptive implementation remains future work.Its residual sampler hierarchically narrows the search, while the decode stage gathers selected keys and values and uses online softmax; clamping inclusion probabilities at 10^-4 introduces a small bias.Source: §4.3 · pp. 7–8
5 Experiments
SSA is evaluated against sparse-attention baselines across multiple models, datasets, context lengths, and generation lengths. It improves benchmark performance and, with optimized kernels, accelerates decoding at long contexts.
5.1 Setup
- Evaluation covers RULER (32K-128K), LongBench (10K-16K), LOFT (32K-128K), and AIME (32K-token generation) across five models from 3B to 27B parameters.Baselines include vAttention, PQCache, and Oracle-top-k; Oracle-top-k uses exact attention scores and is included as an upper-bound reference for deterministic top-k selection.Source: §5.1 · pp. 8–9, §5 · p. 8
5.2 Effectiveness of SSA
- At 50× attention sparsity, SSA outperforms PQ and vAttention(PQ) on RULER-HARD across models and context lengths; it also achieves the lowest non-oracle LOFT attention error and keeps LongBench sparse methods within one point of dense attention.On AIME 2025 with Qwen3.5-9B, SSA scores 59.8 at 50× sparsity, versus 48.3 for oracle top-k, 36.7 for PQ, and 35.8 for vAttention(PQ).Source: §5.2 · pp. 9–10
- For AIME 2025, hybrid heavy-set fractions h ∈[0.2, 0.8] achieve 61.0-64.0, exceeding pure sampling at 40.2-42.5 and pure top-k at 49.4-53.8.At lower sparsity, vAttention(PQ) performs better at 5× and at 10× on Qwen3.5-4B.Source: §5.2 · pp. 9–10
5.3 Decode Efficiency
- At batch size 1, SSA is slower than FlashInfer at 128K, but reaches 1.79× speedup at 1M and 50× sparsity; optimized SSA is faster than FlashInfer across measured batch sizes and sparsity levels.Table 2 reports attention-layer latency at 1M context on one H200; its reference implementation surpasses FlashInfer only at batch size 8.Source: Table 2 · p. 9, §5.3 · p. 10, Figure 3 · p. 10
- At 1M context, optimized SSA reaches 2.96× speedup at batch 16 and 50× sparsity, while the two-level sampler reduces routing overhead by 1.21-3.16× over chunk-only routing.The reported 1M-context sampling latency is 0.30-1.18 ms at h = 0.4; Figure 3 also shows speedup rising with batch size.Source: Figure 3 · p. 10, §5.3 · p. 11
6 Conclusion
SSA unifies deterministic retrieval and sampling, and reduces the accuracy gap to dense attention. It can use any approximate top-k method that supplies proxy scores; adding user-controlled approximation guarantees remains future work.
6 Conclusion
- SSA is a unified framework encompassing top-k, pure sampling, and hybrid sparse attention, and substantially outperforms existing methods while reducing the accuracy gap to dense attention.The framework is compatible with approximate top-k methods that provide proxy scores.Source: §6 · p. 11
A.1 Proof of Lemma 1
The proof fragment bounds token inclusion probabilities and combines the lower bound with an upper bound on a_i.
A.1 Proof of Lemma 1
- The proof uses the decreasing function g(x) = (1 −e−x)/x for x > 0 to establish a two-sided bound on π_i, then combines its lower bound with a_i ≤ ARρp_i.The supplied proof passages do not include the resulting bound's final expression.Source: §A.1 · p. 16
A.2 Proof of Theorem 2
The proof uses a vector-valued bounded-differences inequality to control the error of SSA’s sampled output. It establishes the stated error bound with probability at least 1−δ.
A.2 Proof of Theorem 2
- The proof models the residual as m independent draws, with each residual token’s inclusion probability given by 1−(1−p_i)^m.Heavy tokens are retained deterministically, with inclusion probability one.Source: §A.2 · p. 17
- It defines estimated numerator and denominator quantities for the SSA output, then uses bounded differences to bound their deviations from the corresponding attention quantities.Changing one residual draw can change only the occupancy indicators of the token drawn out and the token drawn in.Source: §A.2 · pp. 16–18
- The proof combines concentration bounds for the numerator and denominator, requiring denominator stability before concluding the output-error bound.A union bound gives simultaneous concentration with probability at least 1−δ.Source: §A.2 · pp. 18–20
- With probability at least 1−δ, the proof establishes the stated bound on the difference between SSA’s output and dense attention’s output.The argument handles the heavy-token contributions exactly and concentrates the residual contributions.Source: §A.2 · pp. 18, 20
A.3 Proof of Corollary 3
The corollary applies the theorem’s bound under the algorithm’s allocation rule and obtains an error guarantee in terms of ε and the value-vector norm.
A.3 Proof of Corollary 3
- Under Algorithm 1’s allocation rule, c < 1 implies κ(c) > 1−e^-1.The proof uses this bound together with the assumed lower bound on m.Source: §A.3 · p. 20
- The proof bounds x below 1 and substitutes that bound into the theorem’s error expression to obtain the ε guarantee.It identifies the numerator of the theorem’s bound as 2x times the value-vector norm.Source: §A.3 · p. 20
- With probability at least 1−δ, the SSA output error is at most ε times the value-vector norm.The proof first shows that the denominator-stability condition holds under the allocation rule.Source: §A.3 · p. 20
A.4 Additional Experiments
The additional experiments report LOFT results for Qwen3.5 models and LongBench results for Qwen3.5-27B. The section also includes tables covering RULER-HARD evaluations and an AIME allocation ablation.
A.4 Additional Experiments
- On LOFT, SSA has the lowest attention error among non-oracle methods at every context length and sparsity level, while Subspan-EM varies by at most 0.04 across methods.The Figure 5 caption describes average Subspan-EM and error across context lengths and sparsity levels.Source: Figure 5 · p. 21
- The additional tables also include RULER-HARD results for Ministral-3-3B and Ministral-3-8B.Table 6 lists the evaluation datasets and subsets; Table 12 covers Ministral-3-8B.Source: Table 6 · p. 22, Table 8 · p. 23, Table 12 · p. 26
- Table 9 reports LongBench performance with Qwen3.5-27B at 10K–16K context; Table 7 covers RULER-HARD with Llama-3.1-8B, and Tables 10 and 11 cover Qwen3.5-4B and Qwen3.5-9B at 50× sparsity.Source: Table 7 · p. 22, Table 9 · p. 24, Table 10 · p. 24
- Table 13 ablates the heavy-hitter fraction h on AIME 2025, ranging from pure importance sampling at h = 0 to pure top-k retrieval at h = 1.Tables 4 and 5 cover theoretical versus practical SSA variants and average RULER-HARD-32K performance across models, respectively.Source: Table 4 · p. 21, Table 5 · p. 21, Table 13 · p. 26