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

In long-context decoding, every generated token reads the whole KV cache from HBM, so the longer the context, the slower attention gets. Sparse attention reads only a small subset of tokens, but at 50× sparsity existing methods fall far short of dense attention: on RULER-HARD 32K with Llama-3.1-8B, dense attention scores 88.76 while PQCache-style top-k gets only 72.51. This paper proposes Stratified Sparse Attention (SSA). The idea is that high-scoring tokens are always kept, the remaining long tail is sampled by proxy score, and each sampled token is up-weighted by its probability of being sampled. At the same 50× sparsity, Llama-3.1-8B's score comes back to 83.67; on AIME 2025 with Qwen3.5-9B, SSA scores 59.8 versus 36.7 for PQCache. The authors also wrote dedicated CUDA kernels, and at 1M context decode attention is 2.96× faster than FlashInfer.

1. The problem: top-k drops the long tail, and uniform sampling wastes budget

With a budget of $k$ tokens, existing methods allocate it in one of three ways. The first is retrieval only: Quest and PQCache pick the $k$ tokens with the highest proxy scores, and Tactic does top-p. This works well when attention is concentrated on a few tokens. Once the distribution is more spread out, though, all attention mass outside the set is lost, and with even a slightly off proxy score, truly important tokens land just outside the boundary. The second is sampling only: MagicPig samples approximately with LSH and applies an importance correction, which covers more ground but has high variance at small budgets. The third is hybrid: vAttention keeps the top-k and samples uniformly from the remaining tokens. It beats the other two, but under aggressive sparsity it is still clearly behind dense attention.

The authors argue that the problem is the word "uniform". The tokens left outside the top-k have very uneven attention weights (softmax probabilities, summing to 1 over all tokens), and uniform sampling spends the same budget on tokens that contribute almost nothing. They first ran a small experiment: for the same budget, they compare the attention output of the sparse approximation with that of full attention and divide the $L_2$ norm of the difference by the norm of the full output, the relative $L_2$ error. Sampling by proxy score gives lower relative error than uniform sampling across budget splits and sparsity levels (figure below). This is the starting point for SSA.

2. Method: stratify, then sample, then correct by probability

Start with the basic form of sampling-based sparse attention. Let $s_i={\mathbf{q}}^{\top }{\mathbf{k}}_i/\sqrt{d}$ and $w_i=\exp (s_i)$. If each token enters the set $\mathcal{S}$ with probability ${\pi }_i$, multiply it by $1/{\pi }_i$:
$$\hat{\mathbf{o}}=\frac{{\sum }_{i\in \mathcal{S}}(w_i/{\pi }_i){\mathbf{v}}_i}{{\sum }_{i\in \mathcal{S}}{w}_i/{\pi }_i}$$
Intuitively, when a token that was unlikely to be sampled does get sampled, it has to speak for a large group of similar tokens that were not, so its weight is scaled up. SSA builds on this form in three steps.

First it decides which tokens are always kept. Apply softmax to the proxy scores ${\tilde{s}}_i$ (in the experiments, PQCache's product-quantization scores) to get $q_i$, then find a $\lambda$ such that ${\sum }_i\min \{1,\lambda q_i\}=k$. Every token with $\lambda q_i\ge 1$ goes into the heavy set $\mathcal{H}$ and is kept directly. The remaining budget $m=k-|\mathcal{H}|$ is left for sampling from the other tokens, the residual set $\mathcal{R}$. This way the size of the heavy set needs no hand-tuning; the proxy distribution decides it. The sharper the distribution, the larger the heavy set.

Then it samples from the residual set. Renormalize $q_i$ over $\mathcal{R}$ to get $p_i$, draw $m$ times with replacement, and deduplicate to get $\mathcal{S}$. The probability that a residual token is drawn at least once is ${\pi }_i=1-(1-p_i{)}^m$, while tokens in the heavy set have ${\pi }_i=1$. Hoot says▾

Hoot says5 tokens, budget 3: the softmax denominator from top-k and from SSAClose ✕

Top-k drops the same tail every time, so its denominator is always too small; SSA multiplies each sampled residual token by 1/π, statistically making up for the unsampled ones: one draw is off, but the weighted denominator averages to the true value.

Top-k, k = 3 (baseline)

proxy q

0.5, 0.2, 0.1, 0.1, 0.1

true w

5, 2, 1, 1, 1, sum = 10

keep top 3

5 + 2 + 1 = 8

result

always 2 short, the tail never counts

The denominator is too small, so the kept tokens are systematically over-weighted.

SSA, same k = 3

λ = 4

λq = 2, 0.8, 0.4…; sum of min = 1 + 0.8 + 1.2 = 3

strata

token 1 goes to the heavy set (π = 1), m = 2

residual p → π

0.4 → 0.64; 0.2 → 0.36

draws tokens 2 and 4

5 + 2 ÷ 0.64 + 1 ÷ 0.36 ≈ 10.9

Token 4's factor 1 ÷ 0.36 ≈ 2.78 is a statistical compensation: on average it makes up the share of unsampled tokens like it.

All numbers are illustrative: the proxy scores q and true weights w are assumed. λ is the factor that scales q so that the min(1, λq) values add up to the budget of 3; π is the probability that a token is drawn. "Not biased on average" means the 1/π-weighted sums (numerator and denominator separately) have the true value as expectation, a general property of weighting by one over the probability of being drawn; numerator and denominator being unbiased separately does not make their ratio (the normalised output) unbiased; the paper itself states an error bound. Source: paper §2: Σ min{1, λq_i} = k fixes the heavy set, m = k − |H|; draw m times with replacement over the residual by p_i and deduplicate, π_i = 1 − (1 − p_i)^m; weights exp(ℓ_i)/π_i, then normalise

Finally it computes exact attention scores ${\ell }_i$ only for $\mathcal{H}\cup \mathcal{S}$, uses weights $\exp ({\ell }_i)/{\pi }_i$, and takes the normalized weighted sum. Tokens drawn more than once are counted once, so the number of tokens actually read never exceeds $k$.

Here is an illustrative example (the numbers are hypothetical). The context is 1 million tokens, sparsity is 50×, and the budget is $k=20,000$. Suppose 8,000 tokens satisfy $\lambda q_i\ge 1$; they are kept directly. The remaining 12,000 draws are made by $p_i$ over roughly 990,000 residual tokens. When a token with a very small $p_i$ is drawn, its $1/{\pi }_i$ is large, and it stands for a large long tail of tokens that each score low but together still hold a fair amount of attention mass. Top-k throws this whole tail away, and uniform sampling spends most of its draws on tokens near zero.

The existing families of methods are all special cases of this framework: with only the heavy set, it is top-k or top-p; with an empty heavy set, it is pure sampling; with a uniform proposal, it is the vAttention-style hybrid.

The authors also prove an error bound that shows how the budget should be split. Let $A_{\mathcal{R}}$ be the true attention mass in the residual set, let ${\bar{a}}_i=a_i/A_{\mathcal{R}}$ be token $i$'s true attention weight renormalized within the residual set, and let $\rho ={\max }_{i\in \mathcal{R}}{\bar{a}}_i/p_i$ measure how far the proposal is from the true residual distribution ($\rho \ge 1$, with equality when they match). Theorem 2 says that with probability at least $1-\delta$, the output error $\Vert \hat{\mathbf{o}}-\mathbf{o}{\Vert }_2$ is of the order $A_{\mathcal{R}}\rho /\sqrt{m}$, with deduplication adding only a constant factor. So keeping more heavy tokens lowers $A_{\mathcal{R}}$, keeping more sampling budget raises $m$, and the two have to be traded off. And once the proxy score underestimates an important residual token, $\rho$ grows and the error rises with it.

The framework does not depend on which approximate retrieval method is used. Any top-k method that produces proxy scores (PQCache, Quest and so on) can serve directly as SSA's proposal, and when the retrieval method improves, SSA improves with it.

3. Implementation: keeping sampling from slowing decode

Long-context decoding is bound by memory bandwidth, and what SSA saves is the traffic of reading KV. To run fast, the actual kernel does not use an adaptive $\lambda$; it fixes the fraction $h$ of the budget given to the heavy set, and the authors report that this matches the theoretical version closely. A decode step has four stages. The first three scan the entire context, but read only a few bytes per token.

The first stage computes proxy scores. Each key is stored as a 16-byte PQ code (16 subspaces, 256 centroids), while an fp16 KV pair at head dimension 256 takes 1 KiB, so scanning the PQ codes is far cheaper than reading KV. The first 128 tokens (sink tokens, where models habitually park a lot of attention) and the latest 128 local tokens are always kept; their scores are masked out in the same kernel so they are not picked again, with no extra pass.

The second stage selects the heavy set with radix top-k: it counts how many tokens fall into each bucket of the score's bits, from the high bits down, and in a few passes finds the threshold for the k-th largest score without sorting a million scores. The returned indices are unsorted, because the estimator above does not depend on order.

The third stage samples the residual. Building a distribution and sampling over a million tokens directly is too slow, so the authors sample hierarchically. Each draw first picks a 1024-token chunk, then a 32-token sub-tile within it, and finally one token among those at most 32. A chunk is chosen with probability equal to the sum of its tokens' $p_i$; within the chosen chunk, a sub-tile is chosen by its share of the chunk's $p_i$, and a token by its share of the sub-tile's. Multiplying the three levels gives each token a final chance of exactly $p_i$, the same as drawing directly over all residual tokens. Several draws for the same query can hit the same token, so each query (one row of the attention matrix) gets its own hash table to remove duplicates, and each unique sample is reweighted by $1/{\pi }_i$. ${\pi }_i$ is clamped from below at $10^{-4}$ to keep individual weights from blowing up, at the cost of a small bias. Random numbers come from a counter-based RNG, which keeps the kernel compatible with CUDA graphs.

The fourth stage gathers the selected K and V from the paged KV cache and computes attention with an online softmax: K and V are read in batches while a running maximum score and sum of exponentials are kept, the accumulated output is rescaled whenever the maximum changes, and one pass yields the normalized result without storing all scores first. To use more parallelism, each query head is split across several thread blocks, and the parts are merged with log-sum-exp. The number of splits matters: at 1M context and batch 1, 64 splits are 1.58–1.71× faster than 32. Going to 128 or 256 splits is slower, because launch and merge overhead outweighs the extra parallelism.

4. Experiments: the biggest gains are at aggressive sparsity, and proxy-score quality is still the bottleneck

There are three baselines: PQCache, vAttention (with PQ scores), and Oracle-top-k. Oracle-top-k computes exact scores before picking the top-k, so it cannot speed anything up; it serves only as a reference for "how good deterministic retrieval can get". All methods keep 128 sink tokens and 128 local tokens. The models are Llama-3.1-8B, Qwen3.5-4B/9B and Ministral-3-3B/8B, plus Qwen3.5-27B on LongBench.

Start with RULER-HARD at 50× sparsity. PQCache, vAttention (PQ) and SSA use the same PQ scores and the same budget and differ only in how the budget is split, so the gaps in the table below can be attributed directly to the allocation.

SSA beats PQCache and vAttention (PQ) in every row. On Ministral-3-8B it retains 94% and 85% of the dense score at 32K and 128K. Note, though, that in the two 128K rows and in the Ministral-3-8B 32K row, Oracle-top-k is still higher than SSA. This shows the proxy score itself is not accurate enough, which is a ceiling SSA cannot get around. On the Qwen3.5 models every sparse method is within 2 points of dense, so there is little for SSA to win.

AIME 2025 requires generating 32K tokens of long reasoning, and here SSA's advantage grows with sparsity. Qwen3.5-9B scores 66.3 dense. At 50× sparsity SSA scores 59.8, while Oracle-top-k gets only 48.3 and PQCache 36.7. At 100× sparsity SSA still has 51.5, while vAttention (PQ) has dropped to 3.5. Conversely, at only 5× sparsity vAttention (PQ) is better, 65.4 to 61.0, so SSA's gains are concentrated in the aggressive-sparsity range.

The heavy fraction $h$ is the share of the budget given to deterministic top-k, with the rest going to sampling. An ablation on the same AIME 2025 task at 10× and 20× sparsity shows that neither part can be dropped: with $h$ between 0.2 and 0.8, scores are all 61.0–64.0. Giving the whole budget to sampling ($h=0$) yields only 40.2–42.5, and giving it all to top-k ($h=1$) yields 49.4–53.8.

LOFT and LongBench show no difference. On LOFT, SSA has the lowest attention error among non-oracle methods in every setting, but the task metric subspan-EM (an answer counts as correct if the gold answer appears as a substring of the output; maximum 1) differs by at most 0.04 across methods. On LongBench (Qwen3.5-27B), every sparse method is within 1 point of dense. In other words, lower error does not necessarily turn into a higher task score; it depends on how sensitive the task is to sparsity.

Speed is measured on a single H200 (16 query heads, 4 KV heads, head dimension 256, fp16 KV). With short contexts SSA is actually slower, because the fixed cost of the first three stages is not amortized: at batch 1 and 128K it runs at only 0.60–0.74× the speed of FlashInfer. The longer the context, the better the trade. At 1M, 50× sparsity is 1.79× faster, 20× is 1.39× faster, and 10× is only 1.07× faster. Larger batches give larger gains, reaching 2.96× at 1M, batch 16 and 50× sparsity.

Kernel optimization is why SSA beats dense even at batch 1. A reference implementation of the same algorithm without the optimizations from Section 3 runs at only 0.37–0.47× FlashInfer and catches up only at batch 8. The optimizations make it 2.95–3.80× faster. For sampling, the comparison is chunk-only routing: after choosing a 1024-token chunk, draw directly from its 1024 tokens without the 32-token sub-block level. Adding the sub-block level makes sampling 1.21–3.16× faster (the gap grows with the number of draws), and at 1M context with $h=0.4$ the sampling itself takes only 0.30–1.18 ms.

5. Limitations and practical value

There are several limitations. The actual kernel uses a fixed fraction $h$, and there is not yet an equally efficient implementation of adaptive allocation. Clamping ${\pi }_i$ introduces a small bias. Sampling is random, so decoding the same request twice does not give identical attention outputs, and reproducibility requires fixing the random seed. Below 128K at batch 1, SSA is slower than FlashInfer. In settings where sparsity barely costs anything to begin with, such as Qwen3.5 or LongBench, SSA has no clear advantage. As a next step, the authors want users to set the approximation error they can accept and have the system guarantee it, but the paper gives no concrete design.

For industry, the scenario where SSA fits is clear: decode services with million-token contexts, large batches, and sparsity pushed beyond 20–50×, such as session histories of long-horizon agents, very long document QA, and long reasoning-chain generation. If your system already has top-k retrieval such as PQCache or Quest, adding SSA is a small change. The retrieval index is reused as is; you only move part of the budget from top-k to sampling by proxy score and add the $1/{\pi }_i$ correction. Before going live, sweep $h$ on your own tasks. In the paper anything from 0.2 to 0.8 works well, while pure top-k and pure sampling are both clearly worse. If you only need 5–10× sparsity, existing top-k or vAttention is probably enough; SSA's gains are mainly in the more aggressive range.