Large language models face a bottleneck in long‑context inference because attending to the growing KV cache is costly. Sparse attention lowers this cost by accessing only a fraction of the cache, yet existing approaches trade deterministic retrieval (capturing high‑scoring tokens) against sampling (improving coverage but increasing variance). Hybrid schemes that combine top‑k retrieval with uniform sampling still show a large gap to dense attention under aggressive sparsity.
We introduce Stratified Sparse Attention (SSA), a unified framework that merges deterministic retrieval with score‑guided sampling. SSA first deterministically attends to the highest‑scoring tokens, then samples the remaining positions using proxy attention scores and corrects for their inclusion probabilities, thereby recovering attention mass beyond the retrieved set. Top‑k, top‑p, and hybrid top‑k+sampling appear as special cases of SSA.
Across the RULER, LongBench, LOFT, and AIME benchmarks and five models ranging from 3B to 27B parameters, SSA consistently outperforms baselines at matched sparsity. On RULER‑HARD with Ministral‑3‑8B, SSA raises accuracy from 81.8% to 84.2% at 32K and from 63.8% to 69.3% at 128K under 50× sparsity, recovering 94% and 85% of dense‑attention performance respectively. With optimized CUDA kernels, SSA achieves up to 2.96× decode speedup over FlashInfer at a 1M context length, showing that accuracy gains translate into practical inference acceleration. The code is released at https://github.com/skylight-org/SSA-Paper-Code.
Review: SSA’s layered sampling and probability correction dramatically narrow the performance gap to dense attention under extreme sparsity while remaining simple to implement and offering notable speedups, making it a strong candidate for long‑context LLM deployment.