Long-context decode is bottlenecked by attending a growing KV cache. Deterministic top-k drops residual mass; sampling adds variance; hybrids still lag dense. Stratified Sparse Attention (SSA) splits the sparse budget between deterministic exploitation and stochastic exploration by strata, closing the sparse–dense gap at extreme sparsity. At 50× sparsity it lifts the strongest practical baseline from 81.8%→84.2% at 32K and 63.8%→69.3% at 128K (94% and 85% of dense); decode vs FlashInfer reaches up to 2.96× at 1M context. At extreme sparsity, stratify to recover residual mass—do not only chase top-k.
1. Problem
Top-k drops residual mass; pure sampling is noisy; top-k+uniform still leaves a dense gap. Need a better split of a fixed budget k.
2. Method
Three stages: allocate budget from proxy scores into deterministic vs sampled parts; sample the residual by strata/probabilities; custom CUDA cuts memory traffic. Uniform sampling and pure top-k are special cases.
3. Core claim
Systematically shrink the sparse–dense gap under extreme sparsity. Speedups count only when quality stays near dense.
4. Experiments
RULER/LongBench/LOFT/AIME across 3B+ models. At 50×: 32K 81.8→84.2, 128K 63.8→69.3. Up to 2.96× decode vs FlashInfer at 1M on H200.
5. Limits
Proxy quality matters; batch size changes speedups; kernel integration is engineering.
6. Practice
If sparse attention is already on, A/B stratified budgets and track dense recovery on RULER/long suites at 50× sparsity.