Stratified Sparse Attention speeds up decoding 2.96x at 1M context

A team of researchers has proposed Stratified Sparse Attention (SSA), a way to access only a small fraction of the KV-cache while losing noticeably less quality than earlier sparse schemes. On RULER-HARD with Ministral-3-8B at 50x attention sparsity, SSA raises accuracy from 81.8% to 84.2% at 32K context and from 63.8% to 69.3% at 128K. That is 94% and 85% of full dense-attention performance. At a context of one million tokens, optimized CUDA kernels speed up decoding by up to 2.96x relative to FlashInfer.

Where sparse attention loses quality

At long context, the inference bottleneck is less the computation itself than reading an ever-growing KV-cache at every decoding step. Sparse attention cuts the cost by touching only part of the cache. This is where the approaches diverge.

The first path is deterministic selection. Top-k or top-p picks the tokens with the highest scores. This reliably catches the most important ones, but it simply throws away all the remaining attention, which is small token by token yet can add up to something significant. The second path is sampling. It covers the tail of the distribution better, but it pays in variance: results are noisier from run to run.

Hybrids glue top-k to uniform sampling. The authors show that at aggressive sparsity, these still leave a large gap to dense attention. Uniform sampling is blind: it spends the budget on tokens that barely affect the answer.

How SSA splits the cache into strata

SSA builds a single scheme from two parts. The model reads high-scoring tokens deterministically. The rest of the cache is neither discarded nor picked at random: tokens from it are sampled with probabilities proportional to proxy attention scores. So that the selected tokens do not distort the picture, SSA applies a correction for their probability of being included in the sample. This restores the attention mass outside the selected set.

The idea is close to classical stratified sampling: cover the important region in full, and estimate the rest from a weighted subsample. The authors stress that SSA is a generalization. Top-k, top-p, and hybrid top-k with sampling turn out to be special cases of it, so the method can be viewed as a general framework for the whole family.

Numbers and speedup

The evaluation ran on four benchmarks (RULER, LongBench, LOFT, and AIME) and five models from 3B to 27B parameters. At equal sparsity, SSA consistently outperforms the baseline methods. Key results:

• RULER-HARD, Ministral-3-8B, 32K: 81.8% → 84.2% against the strongest practical baseline, 94% of dense attention;
• RULER-HARD, Ministral-3-8B, 128K: 63.8% → 69.3%, 85% of dense attention;
• mode: 50x attention sparsity;
• speed: up to 2.96x in decoding relative to FlashInfer at 1M context;
• code is open: github.com/skylight-org/SSA-Paper-Code.

The second part of the result matters too. The method with the best accuracy usually requires more complex operations, and the speedup can vanish. Here, the authors show with optimized CUDA kernels that the quality gain is not eaten up by overhead. The accuracy gain at 128K is larger than at 32K (5.5 pp versus 2.4 pp), meaning the method helps more where the context is longer and sparsity hurts more.

What practitioners get

For teams serving long context, the KV-cache is a direct cost item in GPU memory and bandwidth. If a 50x reduction in cache accesses preserves 85–94% of dense-attention quality on hard tasks, engineers get one more lever between cost and accuracy. This is especially interesting for agentic scenarios and work with large codebases, where context grows fast.

The limitations are visible too. Some of the figures come from a single model, Ministral-3-8B, and 128K still shows a noticeable gap to dense mode: 85% is a good result, but not parity. For production, you will need to test the method on your own tasks and check kernel compatibility with your stack. The open code lowers the barrier: you can reproduce the experiments and compare SSA with current top-k schemes without waiting for an implementation in popular frameworks.