Stratified Sparse Attention: Reducing the Sparse-to-Dense Gap at Extreme Sparsity

Stratified Sparse Attention. Reducing the Sparse-to-Dense Gap at Extreme Sparsity. By Sahil Joshi and seven others, from Rice University, UC Berkeley, and two other institutions. Published on the twentieth of October. Its code repository is worked out: about fifty source files, with a readme and a licence.

Long-context inference in large language models is increasingly limited by the cost of attending to a Key-Value cache that keeps growing. Sparse attention cuts that cost by reading only a fraction of the cache, and existing methods face a trade-off in choosing which fraction. Deterministic retrieval captures the high-scoring tokens but throws away the residual attention mass spread over everything else. Sampling improves coverage of that residual, at the price of higher variance. Hybrid methods combine top-k retrieval with uniform sampling, yet at aggressive sparsity they still sit well short of dense attention.

Stratified Sparse Attention, or SSA, treats the two approaches as strata of one estimate rather than rivals. The high-scoring tokens are attended to deterministically. The rest is sampled, but not uniformly: sampling is guided by scores, and the result is corrected so that what was sampled stands in properly for what was not.

In practice SSA uses proxy attention scores to sample the residual, and corrects for each sampled token's inclusion probability. That correction is what lets it recover attention mass beyond the retrieval set rather than simply ignoring it. The framework is general enough to contain its predecessors: top-k, top-p, and hybrid top-k plus sampling each appear as special cases of it. Optimized CUDA kernels carry it into actual decoding.

Evaluation spans RULER, LongBench, LOFT and AIME, on five models ranging from three billion to twenty-seven billion parameters, and SSA consistently improves over the baselines at matched sparsity. The sharpest numbers come from RULER-HARD with Ministral-3-8B at fifty times attention sparsity.

At thirty-two thousand tokens of context, accuracy over the strongest practical baseline rises from eighty-one point eight percent to eighty-four point two. That recovers ninety-four percent of dense-attention performance.

At one hundred and twenty-eight thousand tokens, accuracy rises from sixty-three point eight percent to sixty-nine point three, recovering eighty-five percent of dense performance. On speed, SSA reaches up to two point nine six times faster decoding than FlashInfer at a context of one million tokens.

The gain comes from refusing to discard the tail. Retrieval alone leaves the residual attention mass on the table, and uniform sampling spends its budget blindly. Sampling that residual by score, and correcting for how likely each token was to be picked, narrows the gap to dense attention exactly where sparsity is most extreme, and the kernels show that the accuracy does not cost speed.