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

Introduction

Large language models (LLMs) are rapidly expanding into long-context workloads, spanning large document collections [23], image analysis [10, 31], software engineering [15, 5], databases [48], video processing [33], and long-horizon agents that continuously accumulate observations, tool outputs, and interaction histories [3, 4]. As context lengths approach millions of tokens, efficient inference becomes increasingly challenging. During autoregressive decoding, each newly generated token attends to an ever-growing key-value (KV) cache, requiring keys and values to be repeatedly transferred from GPU high-bandwidth memory (HBM) onto on-chip SRAM for attention computation [35, 20]. At long contexts, this makes decoding increasingly memory-bandwidth bound: the attention cost per decoding step grows linearly with context length. Sparse attention addresses this bottleneck by attending to only a small fraction of the KV cache [39, 27, 46], reducing both memory traffic and attention latency. However, realizing these gains in practice requires sparse attention to satisfy two competing objectives: faithfully approximate dense attention at aggressive sparsity while remaining efficient enough to translate sparsity into end-to-end inference speedups [29, 25].

Existing sparse-attention methods can be broadly categorized by how they allocate a limited attention budget. First, approximate retrieval methods use proxy attention scores to deterministically select high-scoring tokens, typically through top-$k$ retrieval. Representative methods include PQCache [46], and Quest [39], with recent approaches such as Tactic [50] extending this paradigm to top-$p$ selection. While effective when attention is highly concentrated, deterministic retrieval can miss substantial attention mass when the distribution is diffuse or long-tailed [6, 9]. Second, sampling-based methods such as MagicPig [6] use importance sampling to approximate dense attention, providing broader coverage but introducing sampling variance that can lead to a larger gap from dense attention at limited budgets (see Table 5 in Appendix). Finally, hybrid methods such as vAttention [9] combine deterministic top-$k$ retrieval with uniform sampling [42], preserving high-confidence tokens while extending coverage beyond the top-$k$ set. Although this combination improves the trade-off between exploitation and coverage, a substantial gap to dense attention remains at aggressive sparsity (see Table 5 in Appendix). See § 2 for more details. These observations motivate a fundamental question: Can we design a more effective sampling strategy that substantially narrows the sparse-to-dense accuracy gap?

Model	Dense	PQCache	vAttn(PQ)	Quest	HashAttn	Double Sparsity	MagicPig
Qwen3-4B-Instruct-2507	88.67	85.80	85.67	42.03	–	39.29	18.72
Llama-3.2-3B-Instruct	66.25	39.30	39.59	34.11	–	31.21	17.22
Llama-3.1-8B-Instruct	87.89	68.95	64.42	64.17	49.48	44.62	16.28

Our preliminary studies reveal a key limitation of uniform sampling: while it provides coverage beyond the deterministic top-$k$ set, it ignores the highly non-uniform distribution of residual attention mass⁷. As shown in Figure 1, score-guided importance sampling (the principle that SSA uses) consistently achieves lower relative $L_2$ error than uniform sampling. Motivated by this observation, we introduce Stratified Sparse Attention (SSA), a unified framework that combines deterministic retrieval of high-confidence heavy hitters with score-guided stochastic exploration of the remaining context. Rather than treating all residual tokens equally, SSA uses proxy attention scores to preferentially explore promising regions while preserving stochastic coverage of the long tail. This enables the sparse attention budget to better capture the attention mass that would otherwise be lost under aggressive sparsity. We evaluate SSA across diverse models, including Llama-3.1-8B-Instruct [12], Qwen3.5-4B, Qwen3.5-9B, Ministral-3-3B-Instruct-2517, and Ministral-3-8B-Instruct-2517 [26], and across RULER [16], LongBench [2], LOFT [22], and AIME [28]. Across these evaluations, SSA consistently improves over deterministic, sampling-based, and hybrid sparse-attention baselines, substantially narrowing the gap to dense attention at extreme sparsity. Our contributions can be summarized as follows:

I. Unified Attention Paradigm: We introduce a unified sparse-attention framework that partitions tokens into deterministic and stochastic strata using a generic proxy score. By varying how tokens are retained or sampled, the framework recovers top-$k$, top-$p$, and hybrid top-$k$+sampling as special cases of importance-sampling-inspired sparse attention.

II. Theoretical Guarantees: We provide rigorous bounds for our proposed algorithm, characterizing how closely it approximates dense attention under a limited attention budget.

III. Efficient Kernel Support: We develop efficient CUDA/Triton kernels that realize the computational benefits of our approach in practice, delivering end-to-end decoding speedups over dense attention.

Related Work

In the context of SSA, the extensive prior work can be organized into three broad and comprehensive research directions:

Head-level compression and KV eviction: One line of work reduces KV-cache overhead by sharing or compressing KV heads. Multi-Query Attention (MQA) [37] shares a single KV head across all query heads, while Grouped-Query Attention (GQA) [1] and Multi-Head Latent Attention (MLA) [7] provide more flexible compression to balance efficiency and model quality. A complementary line of work dynamically evicts less important KV entries during inference. Methods such as SnapKV [24], DynamicKV [49], KVZip [19], KeyDiff [30], CurDKV [34], and CapKV [44] use different importance signals and eviction policies to reduce cache size and subsequent attention computation. These approaches are complementary to sparse attention: KV compression reduces the tokens or representations retained across decoding steps, whereas sparse attention dynamically selects which retained tokens to access for each query. The two can therefore be composed to further reduce memory and decoding costs.

Approximate top-$k$/top-$p$ methods: A substantial body of sparse-attention work focuses on finding the tokens with the largest attention scores. Methods such as Quest [39], PQCache [46], SqueezedAttention [14], Adamas [43], SOCKET [18], Sketch-and-Walk [21], RetrievalAttention [27], Loki [38], HashAttention [8], and Double Sparsity [45] employ different proxy scores or retrieval mechanisms to approximate the top-$k$ tokens under dense attention. These methods are effective at recovering high-attention “heavy hitters,” but a fixed top-$k$ selection discards all tokens outside the retrieval set. Consequently, when non-negligible attention mass is distributed across the residual tail, or when approximation error causes relevant tokens to fall just below the top-$k$ boundary, this discarded mass can create a substantial gap from dense attention at aggressive sparsity. Tactic [50] extends this retrieval paradigm from top-$k$ to top-$p$ selection by modeling the decay of attention scores with a power-law distribution and estimating the number of highest-scoring tokens required to satisfy a target attention-mass threshold.

Sampling-based and Hybrid methods: MagicPig [6] addresses the limitations of deterministic top-$k$ attention by using LSH [17] to approximately sample tokens and importance correction to estimate dense attention. vAttention [9] combines deterministic and stochastic selection through a verified hybrid estimator that retains top-$k$ tokens while uniformly sampling from the residual KV cache. KNN-Attention [13] also explores top-$k$ plus sampling via lazy Gumbel sampling, but primarily targets training rather than inference acceleration for existing models. Unlike uniform hybrid sampling, which treats all residual tokens equally, SSA uses approximate attention scores to guide stochastic selection toward tokens more likely to carry residual attention mass (see Figure 1).

Background and Motivation

Sparse Attention

We begin by describing the dense attention computation. For clarity, we consider a single query $\mathbf{q} \in\mathbb{R}^{d}$ attending to a KV cache $\mathbf{K} := (\mathbf{k}_1, \ldots, \mathbf{k}_N)^{\mathsf{T}} \in\mathbb{R}^{N \times d}$ and $\mathbf{V} := (\mathbf{v}_1, \ldots, \mathbf{v}_N)^{\mathsf{T}} \in\mathbb{R}^{N \times d}$.

Full Scaled Dot-Product Attention (SDPA): Let $s_i = \mathbf{q}^{\mathsf{T}}\mathbf{k}_i/\sqrt{d}$ be the similarity score between $\mathbf{q}$ and $\mathbf{k}_i$, and $\mathbf{a} = (a_1, \ldots, a_N)$ with $a_i = w_i/\sum_{j=1}^{N} w_j$ be the corresponding attention weights where $w_i = \exp(s_i)$. We write SDPA output $\mathbf{o} \in\mathbb{R}^{d}$ as

Sparse Attention with deterministic top-$k$ Selection: Let heavy set $\mathcal{H} \subseteq[N]$ denote a deterministic top-$k$ subset from the $N$ cached tokens. Sparse attention over the set $\mathcal{H}$ is defined as $\mathbf{o}_{\mathcal{H}} = \frac{\sum_{i \in\mathcal{H}} w_i\mathbf{v}_i}{\sum_{j \in\mathcal{H}} w_j}$. While deterministic top-$k$ selection concentrates computation on tokens deemed most relevant, tokens outside $\mathcal{H}$ do not contribute to the sparse output.

Sparse Attention with Importance Sampling: Let sampled set $\mathcal{S} \subseteq[N]$ be a non-empty random set of selected tokens. For each token $i$, let $Z_i = 1\{i \in\mathcal{S}\}$ indicate whether it is selected, and let $\pi_i = \Pr(i \in\mathcal{S}) > 0$ be its inclusion probability. Importance-based sparse attention estimates the dense output by weighting each selected token by the inverse of its inclusion probability:

The factors $1/\pi_i$ account for differences in how likely tokens are to be selected. The estimator in eq. (2) is an instance of inverse-probability weighting. We next recall the standard draw-based formulation of importance sampling.

Importance Sampling

Importance sampling [41] is a Monte Carlo technique for estimating a quantity using samples drawn from a proposal distribution. Given a target distribution $p(x)$ and a function $f(x)$, consider the expectation $\mu= \mathbb{E}_{x\sim p}[f(x)] = \sum_x p(x)f(x)$. Instead of sampling directly from $p(x)$, importance sampling draws samples $x_1,\ldots,x_m \sim q(x)$ from a proposal distribution $q(x)$ and rewrites the expectation as $\mu= \mathbb{E}_{x\sim q}\left[\frac{p(x)}{q(x)}f(x)\right]$. This yields the Monte Carlo estimator

The ratio $p(x)/q(x)$ is the importance weight, which corrects for sampling from $q$ rather than $p$. When $q$ places greater probability on samples with large contributions to the expectation, importance sampling can achieve lower variance than uniform sampling with the same sample budget.

Introducing Stratified Sparse Attention (SSA)

As already mentioned, deterministic top-$k$ attention [11, 47] can omit meaningful attention mass outside the selected set. Uniform sampling on residual tokens provides broader coverage, but does not use information about which remaining tokens are more likely to matter. We introduce Stratified Sparse Attention (SSA), a unified framework that retains tokens with high proxy scores and samples the residual tokens according to those scores. SSA divides a fixed attention budget between two well-known complementary objectives: exploitation, by deterministically retaining high-confidence heavy hitters, and exploration, by sampling the residual context according to proxy attention scores. This score-aware stratification preserves heavy hitters while providing meaningful stochastic coverage toward more promising residual tokens. We next formalize the SSA mechanism described Algorithm 1.

Final Algorithm

SSA operates in three stages that divide the sparse-attention budget $k$ between deterministic exploitation and stochastic exploration.

(i) Proxy-based budget allocation (Steps 1-4): SSA converts the proxy scores $\tilde{s}_i$ into probabilities $q_i$ and finds $\lambda$ such that $\sum_{i=1}^{N}\min\{1,\lambda q_i\}=k$. Tokens satisfying $\lambda q_i \ge1$ form the deterministically retained heavy set $\mathcal{H}$. The remaining $m=k-|\mathcal{H}|$ is allocated to importance sampling from the residual set $\mathcal{R}=[N]\setminus\mathcal{H}$. Thus, the proxy distribution determines the split between deterministic allocation and residual sampling automatically; SSA does not require a separate choice of heavy-set size.

(ii) Score-guided residual sampling (Steps 5-10): SSA normalizes the proxy probabilities over $\mathcal{R}$ to obtain proposal probabilities $p_i=q_i/\sum_{j\in\mathcal{R}}q_j$, makes $m$ multinomial draws with replacement, and retains the unique keys in $\mathcal{S}$. Let $\mathcal{K}=\mathcal{H}\cup\mathcal{S}$, and define $Z_i=\mathbf{1}\{i\in\mathcal{K}\}$ to indicate whether token $i$ is selected. Its inclusion probability is $\pi_i=\Pr(Z_i=1)$. For $i\in\mathcal{H}$, $\pi_i=1$. For $i\in\mathcal{R}$, the token is selected if it appears in at least one of the $m$ draws, giving $\pi_i=1-(1-p_i)^m$.

(iii) Exact scoring and correction (Steps 11-15): For each token in $\mathcal{K}$, SSA computes its exact score $\ell_i=\mathbf{q}^{\mathsf{T}}\mathbf{k}_i/\sqrt{d}$ and corrected weight $\hat{w}_i=\exp(\ell_i)/\pi_i$. It returns the normalized weighted sum $\hat{o}$. Since repeated draws are retained only once, $|\mathcal{K}|\le k$.

Algorithm 1 Stratified Sparse Attention (SSA)

1: Input: Proxy scores $\{\tilde{s}_i\}_{i=1}^{N}$, query $\mathbf{q}$, keys $\{\mathbf{k}_i\}_{i=1}^{N}$, values $\{\mathbf{v}_i\}_{i=1}^{N}$, total budget $k$

2: $\mathbf{q}_i \leftarrow\frac{\exp(\tilde{s}_i)}{\sum_{j=1}^{N}\exp(\tilde{s}_j)}$

3: Find $\lambda> 0$ s.t. $\sum_{i=1}^{N}\min(1,\lambda q_i)=k$

4: Define heavy set $\mathcal{H}=\{i:\lambda q_i\ge1\}$, residual set $\mathcal{R}=[N]\setminus\mathcal{H}$, heavy size $h=|\mathcal{H}|$, sampling budget $m=k-h$

5: Normalize the proposal over $\mathcal{R}$: $p_i=\frac{q_i}{\sum_{j\in\mathcal{R}}q_j}$ for $i\in\mathcal{R}$

6: Draw with replacement: $J_1,J_2,\ldots,J_m\sim p$

7: $\mathcal{S}\leftarrow\operatorname{unique}\{J_1,J_2,\ldots,J_m\}$

8: $\pi_i\leftarrow1-(1-p_i)^m$ for $i\in\mathcal{S}$

9: $\mathcal{K}\leftarrow\mathcal{H}\cup\mathcal{S}$

10: $\pi_i\leftarrow1$ for $i\in\mathcal{H}$

11: for $i\in\mathcal{K}$ do

12: $\ell_i\leftarrow\mathbf{q}^{\top}\mathbf{k}_i/\sqrt{d}$

13: $\hat{w}_i\leftarrow\exp(\ell_i)/\pi_i$

14: end for

15: return $\hat{\mathbf{o}}\leftarrow\frac{\sum_{i\in\mathcal{K}}\hat{w}_i\mathbf{v}_i}{\sum_{i\in\mathcal{K}}\hat{w}_i}$

Theoretical Insights

We now characterize how accurately SSA approximates dense attention. Our analysis of Algorithm 1 is designed for the budget parameter $k\in\{1,\ldots,N-1\}$. When $k=N$, every token is retained (see step 3), so SSA trivially returns dense attention exactly. The residual set is then empty and requires no sampling analysis. For $k<N$, the residual set is nonempty i.e., $m\ge1$. Let $a_i$ be the dense attention weights from § 3, and $\lVert V\rVert_2$ denote the spectral norm of the value matrix. In addition, we introduce the following quantities for our analysis,

where $A_{\mathcal{R}}$ is the true attention mass left for sampling, while $\rho$ measures how well the residual proposal covers that mass. For finite proxy scores, $q_i>0$ and hence $p_i=q_i/\sum_{j\in\mathcal{R}}q_j>0$ for every $i\in\mathcal{R}$. Since $\sum_{i\in\mathcal{R}}p_i(\bar{a}_i/p_i)=\sum_{i\in\mathcal{R}}\bar{a}_i=1$, the largest ratio must satisfy $\rho\ge1$. The equality holds precisely when every ratio is equal to one, i.e., $p_i=\bar{a}_i$ for all $i\in\mathcal{R}$.

The budget-allocation rule also limits the effect of duplicate residual draws. Indeed, steps 3–4 of Algorithm 1 give $k=h+\lambda\sum_{i\in\mathcal{R}}q_i$ as each heavy token contributes 1 to the budget equation, and each residual token contributes $\lambda q_i<1$. Since $m=k-h$, it follows that $m=\lambda\sum_{i\in\mathcal{R}}q_i$. Combining this with $p_i=q_i/\sum_{j\in\mathcal{R}}q_j$ gives $mp_i=\lambda q_i<1$ for every $i\in\mathcal{R}$. Here $mp_i$ is the expected number of times token $i$ appears among the $m$ draws. The largest expected multiplicity among the residual tokens is therefore defined as

A small value of $c$ corresponds to a diffuse proposal with few repetitions, whereas a value closer to one indicates a more concentrated proposal and a greater possibility of repeated draws.

Now, to use the proxy-quality parameter $\rho$ in the error bound, we need to relate the proposal probability $p_i$ to the inclusion probability $\pi_i$ used by Algorithm 1. A residual token is retained if it appears in at least one of the $m$ draws, so $\pi_i = 1 - (1 - p_i)^m$. To account for duplicate draws, we define

The following lemma quantifies their effect on the inclusion probabilities.

Now, we are ready to use the lemma to bound SSA’s approximation error.

Since $\kappa(c) > 1-e^{-1}$, deduplication contributes only a universal constant factor to the approximation and sample-complexity bounds. The following corollary gives a direct $(\varepsilon,\delta)$ guarantee and a sufficient sample size $m = O\left(A_{\mathcal{R}}^2\rho^2\varepsilon^{-2}\log\frac{1}{\delta}\right)$ that automatically ensures $\epsilon_D(\delta) < 1$.

Remark 4. Note that the approximation error in Theorem 2 scales as $A_{\mathcal{R}}\rho/\sqrt{m}$. The residual mass $A_{\mathcal{R}}$ decreases when the retained heavy tokens capture more of the true attention mass, while $\rho$ increases when the proxy underestimates an important residual token. Increasing the residual budget $m$ reduces the stochastic error at the standard Monte Carlo rate $m^{-1/2}$. For a fixed total budget, SSA therefore balances the reduction in $A_{\mathcal{R}}$ obtained from retaining more heavy tokens against the larger sampling budget obtained from retaining fewer of them.

Remark 5. In both Theorem 2 and Corollary 3, we express the bounds using $\lVert V\rVert_2$. The proofs only require a bound on the individual value-vector norms, so the same results hold with $\max_{i \in[N]}\lVert v_i\rVert_2$ in its place, without changing the constants.

Kernel Design

At long contexts, decode attention is memory-bandwidth bound because each step reads the entire KV cache from HBM [35, 20]. SSA reduces this memory traffic by selecting a sparse subset of tokens before loading their KV states. While Algorithm 1 determines the division between the heavy and residual strata adaptively, our optimized GPU implementation uses a fixed heavy fraction. Developing an equally efficient implementation of the adaptive algorithm remains future work. As shown in Table 4, the practical fixed-budget variant closely matches the performance of the theoretical formulation. The residual proposal, inclusion-probability correction, and resulting attention estimator remain unchanged. Our SSA decode step consists of four stages, where only the first three scan the full context using a few bytes per token.

Model	Sparsity	SSA-Theory	SSA-Practical
Llama-3.1-8B	50×	81	83
Llama-3.1-8B	20×	88	88
Ministral-3-8B	50×	82	84
Ministral-3-8B	20×	89	88

Proxy scoring: Each key is represented by a 16-byte PQ code (16 subspaces, 256 centroids) [46], compared to 1 KiB for its fp16 key-value pair at head dimension 256. For each query, a lookup table computes the proxy score $\tilde{s}_i$ for every token. Scores are scaled and masked for sink and local tokens within the same kernel, avoiding an additional pass.

Heavy-set selection: We select $\mathcal{H}$ using radix top-$k$ [36], which returns an unsorted index set since the estimator in eq. (2) is invariant to the ordering of selected tokens.

Residual sampling: Direct sampling from the full residual distribution becomes expensive at long context lengths, as constructing and sampling from the distribution requires processing probabilities across the entire residual KV cache. To reduce this overhead, we use a hierarchical sampler that progressively narrows the search space. Each draw first selects a 1024-token chunk, then a 32-token sub-tile, and finally scans at most 32 tokens to select a key. We remove duplicate samples with a per-row hash table and reweight each unique sample by $1/\pi_i$. In practice, we clamp $\pi_i$ at $10^{-4}$ to avoid excessively large weights, at the cost of introducing a small bias. Finally, we use a counter-based random number generator to obtain fresh samples across replays while retaining compatibility with CUDA graph capture.

Sparse decode: The final stage gathers the keys and values in $\mathcal{H} \cup\mathcal{S}$ from the paged KV cache and evaluates eq. (2) using an online softmax. To exploit more parallelism, we divide each query head across multiple thread blocks and combine their partial softmax results using a numerically stable log-sum-exp reduction. At 1M context and batch size 1, using 64 splits is 1.58–1.71× faster than using 32 splits across sparsity levels. Increasing the number of splits to 128 or 256, however, reduces performance because the additional parallelism no longer offsets the increased launch and merge overhead.

Experiments

We conduct an extensive evaluation of SSA against Oracle-top-$k$, deterministic top-$k$, pure sampling, and hybrid baselines across multiple datasets, models, context lengths, and generation lengths. In addition, note that we slightly overload the notation $h$. In Algorithm 1 and the theoretical analysis, $h = |\mathcal{H}|$ denotes the number of heavy tokens, whereas in the implementation and experimental discussion, $h \in[0, 1]$ denotes the fraction of the total budget assigned to the heavy set. We describe the evaluation setup below.

Setup

Baselines: We compare SSA against three representative sparse-attention baselines: vAttention [9], PQCache [46], and Oracle-top-$k$. Oracle-top-$k$ performs exact attention-score computation before selecting the top-$k$ tokens and is therefore impractical for acceleration; we include it as an upper-bound reference for deterministic top-$k$ selection. For all methods, we use 128 sink tokens and 128 local tokens, along with the hyperparameters recommended by the respective authors. Additional comparisons with MagicPig [6], HashAttention [8], Double Sparsity [45], and Quest [39] are provided in Table 5 in the Appendix. These results motivate our use of PQCache, the strongest approximate top-$k baseline in this comparison, to construct SSA’s proposal distribution throughout our experiments.

Datasets and Models: We evaluate SSA on four benchmark suites: RULER (32K-128K) [16], LongBench (10K-16K) [2], LOFT (32K-128K) [22], and AIME (32K token generation) [28]. Our evaluation spans five models: Llama-3.1-8B-Instruct [12], Qwen3.5-4B and Qwen3.5-9B [40], and Ministral-3-3B-Instruct-2517 and Ministral-3-8B-Instruct-2517 [26]. Dataset subsets and additional evaluation details are provided in Table 6 in Appendix.

Dataset	# Samples	Subsets
RULER-HARD	100	fwe, qa1, qa2, niah_multikey_2, niah_multikey_3, vt
LongBench	200	NrtvQA, Qasper, MF-en, HotpotQA, 2WikiMQA, Musique, DuReader, GovRep, QMSum, MultiNews, VCSum, TREC, TriviaQA, SAMSum, PCount, PR-en, LCC, RBP
LOFT	100	HotpotQA, Musique, NQ, Qampari, Quest
AIME	16	–

Effectiveness of SSA

RULER-HARD: Table 1 reports average RULER-HARD scores at $50 \times$ attention sparsity. SSA consistently outperforms both PQ and vAttention(PQ) across all models and context lengths. Since these methods use the same PQ scores and attention budget, the comparison highlights the benefit of allocating part of the budget to score-guided sampling rather than exclusively to top-$k$ retrieval or uniform sampling. The gains are particularly pronounced when deterministic top-$k$ struggles. At 32K, SSA improves Llama-3.1-8B from 72.51 (PQ) to 83.67, approaching the dense score of 88.76. Similarly, on Ministral-3-3B, SSA improves the score from 70.81 to 82.05, compared with 86.69 for dense attention. Oracle-top-$k$, which selects tokens using exact attention scores, remains stronger at 128K for Llama-3.1-8B and Ministral-3-3B, and across all context lengths for Ministral-3-8B. On Qwen3.5-4B and Qwen3.5-9B, all sparse methods remain within two points of dense attention.

Model	Method	32K	64K	128K
Llama-3.1-8B	Dense	88.76	86.52	71.26
Llama-3.1-8B	PQ	72.51	67.53	39.73
Llama-3.1-8B	vAttn (PQ)	66.71	66.22	39.45
Llama-3.1-8B	Oracle-top-$k$	78.80	71.91	56.10
Llama-3.1-8B	SSA	83.67	77.09	47.53
Ministral-3-3B	Dense	86.69	82.19	81.49
Ministral-3-3B	PQ	70.81	69.88	66.48
Ministral-3-3B	vAttn (PQ)	69.26	65.77	63.06
Ministral-3-3B	Oracle-top-$k$	72.81	75.09	71.87
Ministral-3-3B	SSA	82.05	75.70	69.65
Ministral-3-8B	Dense	89.33	86.13	81.45
Ministral-3-8B	PQ	81.83	78.02	62.99
Ministral-3-8B	vAttn (PQ)	78.69	75.19	63.78
Ministral-3-8B	Oracle-top-$k$	86.61	83.41	75.83
Ministral-3-8B	SSA	84.17	79.91	69.25
Qwen3.5-4B	Dense	91.95	91.11	89.45
Qwen3.5-4B	PQ	91.22	90.72	88.72
Qwen3.5-4B	vAttn (PQ)	90.56	90.61	89.28
Qwen3.5-4B	Oracle-top-$k$	91.39	90.39	89.39
Qwen3.5-4B	SSA	91.45	90.83	89.50
Qwen3.5-9B	Dense	93.45	92.22	91.28
Qwen3.5-9B	PQ	92.00	91.08	90.95
Qwen3.5-9B	vAttn (PQ)	91.67	91.30	90.61
Qwen3.5-9B	Oracle-top-$k$	92.56	92.28	91.39
Qwen3.5-9B	SSA	93.33	92.36	91.45

AIME 2025: Table 3 reports avg@16 accuracy for long-form generation. SSA remains robust as sparsity increases, with particularly large gains in the aggressive-sparsity regime. On Qwen3.5-9B at $50 \times$ sparsity, SSA achieves 59.8, compared with 48.3 for oracle top-$k$, 36.7 for PQ, and 35.8 for vAttention(PQ). At $100 \times$ sparsity, SSA retains an accuracy of 51.5, while vAttention(PQ) drops to 3.5. At lower sparsity, vAttention(PQ) performs better at $5 \times$ and at $10 \times$ on Qwen3.5-4B. We further study the effect of the heavy fraction $h$ in Table 13. Hybrid allocations with $h \in[0.2, 0.8]$ achieve 61.0-64.0, substantially outperforming both pure sampling ($h = 0$), which achieves 40.2-42.5, and pure top-$k$

Sparsity	Qwen3.5-9B Oracle-top-$k$	PQ	vAttn(PQ)	SSA	Qwen3.5-4B Oracle-top-$k$	PQ	vAttn(PQ)	SSA	Method	50$\times$	20$\times$
Dense		66.3				54.2			Dense	41.56	41.56
5$\times$	63.8	63.5	65.4	61.0	49.8	49.4	54.4	52.7	PQ	40.60	40.98
10$\times$	59.4	53.8	62.7	64.0	47.9	41.5	51.9	51.0	vAttn (PQ)	40.39	41.47
20$\times$	55.4	49.4	58.3	63.5	42.7	33.8	50.2	53.1	SSA	40.92	41.56
50$\times$	48.3	36.7	35.8	59.8	36.5	25.4	27.5	50.0	Oracle-top-$k$	41.25	41.40
100$\times$	42.5	28.1	3.5	51.5	31.9	13.1	1.7	42.5

A fraction $h$ of the sparse attention budget is allocated to deterministic heavy hitters, while the remaining $1-h$ is allocated to importance sampling. Thus, $h=0$ corresponds to pure importance sampling and $h=1$ corresponds to PQCache (pure top-$k$ retrieval).

Heavy Fraction $h$	20×	10×
0.0	40.2	42.5
0.2	61.0	63.8
0.4	63.5	64.0
0.6	62.7	62.7
0.8	63.8	63.1
1.0 (PQ)	49.4	53.8

($h = 1$), which achieves 49.4-53.8. These results demonstrate the benefit of combining deterministic retrieval with score-guided importance sampling, particularly under aggressive sparsity.

LOFT and LongBench: On LOFT with Qwen3.5-9B and Qwen3.5-4B (see Figure 2), SSA has the lowest attention error among the non-oracle methods at every context length and sparsity level, while subspan-EM varies by at most 0.04 across methods (see Figure 5). On LongBench with Qwen3.5-27B (Table 3), all sparse methods stay within one point of dense attention at both 50$\times$ and 20$\times$ sparsity.

Decode Efficiency

Sparse attention accelerates decoding only when its selection overhead is smaller than the memory traffic it saves. We evaluate this trade-off for SSA across context lengths and batch sizes, and separately measure the sampling overhead absent in top-$k$ methods.

Efficiency analysis: SSA becomes increasingly advantageous as context length and batch size grow. At batch size 1, SSA is slower than FlashInfer at 128K, achieving approximately 0.60-0.74$\times$ its speed. As context length increases, this overhead is increasingly amortized, and SSA reaches 1.79$\times$, 1.39$\times$, and 1.07$\times$ speedups at 1M for 50$\times$, 20$\times$, and 10$\times$ sparsity, respectively (see Figure 3a). The gains further increase with batch size, reaching $2.96 \times$ at batch 16 for $50 \times$ sparsity (see Figure 3b); our kernel optimizations are particularly important at batch 1, improving SSA by $2.95$-$3.80 \times$ over the reference implementation (Table 2). Finally, although residual sampling introduces overhead absent from top-$k$ methods, our two-level sampler reduces this cost by $1.21$-$3.16 \times$ over chunk-only routing as the number of draws $m$ increases (see Figure 3c), keeping sampling latency at $0.30$-$1.18$ ms at 1M context and $h = 0.4$ across the evaluated sparsity levels.

Sparsity	Batch	Latency (ms)			SSA speedup
		FlashInfer	Reference	SSA	vs. FlashInfer	vs. reference
$50 \times$	1	1.41	2.99	0.79	1.79×	3.80×
	8	11.90	7.18	4.27	2.78×	1.68×
	16	24.47	12.57	8.26	2.96×	1.52×
$20 \times$	1	1.41	3.34	1.01	1.40×	3.32×
	8	11.91	8.74	5.76	2.07×	1.52×
	16	24.53	15.51	11.32	2.17×	1.37×
$10 \times$	1	1.42	3.83	1.30	1.10×	2.95×
	8	12.01	11.06	7.91	1.52×	1.40×
	16	24.55	19.82	15.81	1.55×	1.25×

Conclusion

We introduce Stratified Sparse Attention (SSA), a unified framework encompassing top-$k$, pure sampling, and hybrid sparse attention while substantially outperforming existing methods and reducing the accuracy gap to dense attention. SSA is compatible with any approximate top-$k$ method that provides proxy scores, allowing it to directly benefit from advances in approximate retrieval. An important direction for future work is to make SSA verifiable by incorporating user-controlled guarantees on approximation quality.

Acknowledgments

The work was primarily supported by Rice Ken Kennedy Institute (K2I) Generative AI Cluster Funding. The Sky Computing Lab is supported by gifts from Accenture, Algorithmic SuperIntelligence Labs, Amazon, AMD, Anyscale, Broadcom, cmpnd, Google, IBM, Intel, Intesa Sanpaolo, Lambda, Lightspeed, Mirendil, NVIDIA, Samsung SDS, and VESSL.

Appendix

Proof of Lemma 1

Proof. Token $i$ is included if it appears in at least one of the $m$ independent draws, so

By the union bound,

For the lower bound, $1-p_i \le e^{-p_i}$ gives

The function $g(x)=(1-e^{-x})/x$ is decreasing for $x>0$. Since $mp_i\le c$,

This proves the two-sided bound on $\pi_i$.

Finally, the definition of $\rho$ gives $a_i=AR\bar a_i\le AR\rho p_i$. Combining this with the lower bound on $\pi_i$ yields

Proof of Theorem 2

We first establish the following dimension-free bounded-differences inequality used in the proof.

Proof. Reveal the variables one at a time and define

Here

The sequence $(M_\ell)_{\ell=0}^{m}$ is a vector-valued martingale.

We first bound its increments. Fix the first $\ell-1$ variables and let

Averaging over the same future variables preserves the bounded-differences condition:

Let $X'_{\ell}$ be an independent copy of $X_{\ell}$, independent of all the original variables. Since $M_{\ell}=g_{\ell}(X_{\ell})$ and $M_{\ell-1}$ averages $g_{\ell}$ over the $\ell^{\mathrm{th}}$ variable, Jensen’s inequality gives

The Euclidean specialization of the martingale inequality in [32 Theorem 3.5], applied to $(M_{\ell}-M_{0})_{\ell=0}^{m}$, therefore yields

Substituting $M_{m}$ and $M_{0}$ proves the claim.

Now, we are ready to prove Theorem 2.

Proof of Theorem 2. We consider the sparse regime with an integer budget $1\le k<N$ and finite attention and proxy scores. The case $k=N$ can be handled by evaluating dense attention exactly. Recall that $\mathcal{H}$ and $\mathcal{R}=[N]\setminus\mathcal{H}$ be the heavy and residual sets produced by SSA, and let $h=|\mathcal{H}|$ and $m=k-h$. As shown in § 4.2, $\mathcal{R}\ne\varnothing$ and $m\ge1$. Since the proxy scores are finite, $p_i>0$ for every $i\in\mathcal{R}$.

Let $J_{1},\ldots,J_{m}$ denote the independent residual draws:

For each residual token, define its occupancy indicator

We know that its inclusion probability is given by $\pi_i=\Pr(Z_i=1)=1-(1-p_i)^m$, so $\mathbb{E}[Z_i]=\pi_i$. Heavy tokens are retained deterministically and have inclusion probability one.

Note that although the residual indicators $\{Z_i\}_{i\in\mathcal{R}}$ are dependent, they are deterministic functions of the independent draws $J_1,\ldots,J_m$. We therefore apply Lemma 6 to these draws.

Write

Define the estimated denominator and numerator by

The SSA output is therefore

Since $m \ge1$ and all weights are positive, $\widehat{D} > 0$ for every realization of the residual draws.

If $\lVert V\rVert_{2} = 0$, all value vectors vanish, and $\widehat{o} = o = 0$, so the theorem is immediate. Therefore, we assume $\lVert V\rVert_{2} > 0$.

By Lemma 1,

For brevity, define

Thus $a_i/\pi_i \le\beta$ for every $i \in\mathcal{R}$.

Next, we define

The contributions from $\mathcal{H}$ are evaluated exactly and cancel in $\widehat{D} - D$. Consequently,

Since $\mathbb{E}[Z_i] = \pi_i$,

Consider two draw sequences $\mathbf{J} = (J_1,\ldots,J_m)$ and $\mathbf{J}'$ differing only in coordinate $\ell$, where the draw changes from $r$ to $s$. Only $Z_r$ and $Z_s$ can change, and each changes by at most one. Therefore,

Applying Lemma 6 with $r = 1$ and $b_\ell= 2\beta$ gives, for every $t > 0$,

Taking $t = 2\sqrt{2m}\,\beta\sqrt{u}$, we obtain, for every $u > 0$,

Substituting eq. (6) yields

Now, we define

Because $\sum_{i=1}^{N} a_i = 1$ and $o = \sum_{i=1}^{N} a_i v_i$, we have

It follows that

Again using $\mathbb{E}[Z_i]=\pi_i$, we obtain

Let $\mathbf{e}_i$ denote the $i^{\mathrm{th}}$ standard basis vector of $\mathbb{R}^{N}$. Since the value vectors are the rows of $\mathbf{V}$, for every $i\in[N]$,

Moreover, $\mathbf{o}$ is a convex combination of the value vectors, so

Therefore,

Now consider draw sequences $\mathbf{J}$ and $\mathbf{J}'$ that differ only in one draw, changed from $r$ to $s$. As before, only $Z_r$ and $Z_s$ can change. Hence,

Applying Lemma 6 with $r=d$ and $b_\ell=4\beta\|\mathbf{V}\|_2$ gives, for every $t>0$,

Taking $t=4\sqrt{2m}\,\beta\sqrt{u}\,\|\mathbf{V}\|_2$ gives, for every $u>0$,

Equivalently,

Now, fix the failure probability $\delta\in(0,1)$ and set

Then $2e^{-u}=\delta/2$. By eq. (11), with probability at least $1-\delta/2$,

Similarly, by eq. (17), with probability at least $1-\delta/2$,

A union bound shows that both inequalities hold simultaneously with probability at least $1-\delta$. By the definition of $\Delta_D$,

Consequently,

On the simultaneous concentration event, the assumption $\epsilon_D(\delta)<1$ implies

Therefore,

This proves eq. (3) with probability at least $1-\delta$.

Proof of Corollary 3

Proof. Let

The assumed lower bound on $m$ gives

Under the Algorithm 1 allocation rule, $c<1$, so $\kappa(c)>1-e^{-1}$. Since

we have

Thus the denominator-stability condition in Theorem 2 holds. The numerator of the bound in Theorem 2 equals $2x\left\lVert\mathbf{V}\right\rVert_2$. Therefore, with probability at least $1-\delta$,

Additional Experiments

Sparsity	Context	Method	FWE	VT	NM2	NM3	QA1	QA2	Avg.
$50\times$	32K	Dense	93.33	99.2	100	100	88	52	88.76
$50\times$	32K	PQ	77.67	95.4	67	71	76	48	72.51
$50\times$	32K	vAttn (PQ)	83.67	94.6	52	58	66	46	66.71
$50\times$	32K	Oracle-top-$k$	81.00	97.8	87	76	79	52	78.80
$50\times$	32K	SSA	91.00	98.0	83	92	85	53	83.67
$50\times$	64K	Dense	86.33	98.8	98	99	84	53	86.52
$50\times$	64K	PQ	80.00	92.2	65	50	71	47	67.53
$50\times$	64K	vAttn (PQ)	83.33	94.0	40	68	67	45	66.22
$50\times$	64K	Oracle-top-$k$	79.67	95.8	78	56	77	45	71.91
$50\times$	64K	SSA	85.33	96.2	67	84	77	53	77.09
$50\times$	128K	Dense	73.33	66.2	91	67	83	47	71.26
$50\times$	128K	PQ	66.00	39.4	27	15	57	34	39.73
$50\times$	128K	vAttn (PQ)	71.67	55.0	21	20	41	28	39.45
$50\times$	128K	Oracle-top-$k$	61.00	66.6	71	32	66	40	56.10
$50\times$	128K	SSA	73.00	54.2	30	30	59	39	47.53
$20\times$	32K	PQ	87.33	98.6	97	93	84	52	85.32
$20\times$	32K	vAttn (PQ)	93.33	99.4	97	99	84	52	87.46
$20\times$	32K	Oracle-top-$k$	90.33	98.8	100	99	84	52	87.36
$20\times$	32K	SSA	93.67	99.0	97	100	87	54	88.45
$20\times$	64K	PQ	86.33	96.2	91	91	77	55	82.76
$20\times$	64K	vAttn (PQ)	86.00	98.4	87	98	79	57	84.23
$20\times$	64K	Oracle-top-$k$	85.33	98.0	94	95	81	52	84.22
$20\times$	64K	SSA	85.33	98.6	94	99	81	54	85.32
$20\times$	128K	PQ	69.00	65.8	66	46	68	34	58.13
$20\times$	128K	vAttn (PQ)	73.33	63.6	61	51	69	39	59.49
$20\times$	128K	Oracle-top-$k$	67.00	63.2	85	51	72	42	63.37
$20\times$	128K	SSA	75.00	69.0	70	63	76	44	66.17

Sparsity	Context	Method	FWE	VT	NM2	NM3	QA1	QA2	Avg.
$50\times$	32K	Dense	96.33	91.8	98	99	79	56	86.69
$50\times$	32K	PQ	92.67	60.2	91	56	73	52	70.81
$50\times$	32K	vAttn (PQ)	90.33	72.2	67	77	65	44	69.26
$50\times$	32K	Oracle-top-$k$	93.67	62.2	99	55	75	52	72.81
$50\times$	32K	SSA	95.67	88.6	92	90	72	54	82.05
$50\times$	64K	Dense	74.33	99.8	93	96	75	55	82.19
$50\times$	64K	PQ	72.67	94.6	77	51	69	55	69.88
$50\times$	64K	vAttn (PQ)	76.00	95.6	54	75	52	42	65.77
$50\times$	64K	Oracle-top-$k$	75.33	98.2	89	68	68	52	75.09
$50\times$	64K	SSA	75.00	99.2	76	87	67	50	75.70
$50\times$	128K	Dense	92.33	98.6	95	87	66	50	81.49
$50\times$	128K	PQ	77.67	99.2	66	43	66	47	66.48
$50\times$	128K	vAttn (PQ)	89.33	98.0	54	42	54	41	63.06
$50\times$	128K	Oracle-top-$k$	88.00	98.2	79	51	66	49	71.87
$50\times$	128K	SSA	91.67	98.2	60	67	60	41	69.65
$20\times$	32K	PQ	93.33	67.6	98	86	79	54	79.66
$20\times$	32K	vAttn (PQ)	95.33	90.8	97	98	76	52	84.86
$20\times$	32K	Oracle-top-$k$	96.00	73.0	99	93	75	52	81.33
$20\times$	32K	SSA	95.67	94.2	99	97	79	56	86.81
$20\times$	64K	PQCache	77.00	98.4	87	84	70	56	78.73
$20\times$	64K	vAttn (PQ)	73.00	99.8	86	96	73	53	80.13
$20\times$	64K	Oracle	77.00	99.2	93	91	67	52	79.87
$20\times$	64K	SSA	75.33	99.8	86	96	73	56	81.02
$20\times$	128K	PQ	88.33	99.2	71	65	70	50	73.92
$20\times$	128K	vAttn (PQ)	89.33	98.4	59	77	65	44	72.12
$20\times$	128K	Oracle-top-$k$	91.00	98.4	85	75	68	50	77.90
$20\times$	128K	SSA	90.33	98.6	75	81	67	46	76.32

Sparsity	Method	NrtvQA	Qasper	MF-en	MF-zh	HotpotQA	2WikiMQA	Musique	DuReader	GovRep	QMSum
50×	Dense	20.49	29.64	38.94	42.97	34.83	34.64	26.69	25.92	31.91	21.96
50×	PQ	20.13	28.69	38.60	42.01	33.27	31.92	20.61	25.54	31.92	21.64
50×	vAttn(PQ)	21.97	25.33	39.38	43.53	32.00	30.44	21.45	26.91	30.68	22.42
50×	SSA	20.42	27.51	38.87	42.13	34.64	32.79	24.88	24.82	31.96	22.72
50×	Oracle-top-$k$	20.51	28.97	38.22	42.12	34.79	33.77	22.50	25.76	32.86	22.52
20×	PQ	20.61	29.55	38.41	41.76	34.51	33.38	23.90	25.74	31.63	22.31
20×	vAttn(PQ)	21.55	27.17	40.14	43.67	33.39	34.18	25.65	26.70	32.20	22.27
20×	SSA	20.90	28.80	39.73	42.91	34.41	33.97	25.24	25.93	31.36	22.58
20×	Oracle-top-$k$	20.74	29.28	38.68	42.87	34.90	34.57	23.46	25.82	32.77	21.97

Sparsity	Method	MultiNews	VCSUM	TREC	TriviaQA	SAMSum	PCount	PRe-en	PRe-zh	LCC	RB-P	Avg.
50×	Dense	23.33	14.50	72.50	89.65	45.63	5.16	88.12	57.33	67.96	59.12	41.56
50×	PQ	24.37	14.50	71.25	89.68	45.92	6.37	87.40	56.60	66.65	54.97	40.60
50×	vAttn(PQ)	23.17	14.62	71.25	89.04	43.76	6.85	85.25	59.21	66.46	53.98	40.39
50×	SSA	23.08	14.35	72.50	89.77	44.32	5.16	86.25	59.69	64.55	58.01	40.92
50×	Oracle-top-$k$	24.62	14.66	72.50	89.66	46.91	5.28	87.81	53.90	68.56	59.06	41.25
20×	PQ	24.20	14.79	70.00	89.46	45.83	5.10	89.06	54.21	67.87	57.24	40.98
20×	vAttn(PQ)	23.41	15.05	75.00	90.52	45.12	4.68	86.88	57.23	66.02	58.56	41.47
20×	SSA	24.01	14.67	72.50	90.18	46.93	4.76	88.12	56.81	66.85	60.47	41.56
20×	Oracle-top-$k$	24.75	14.29	71.25	89.66	46.25	5.26	87.71	55.04	68.84	59.90	41.40

Context	Method	FWE	VT	NM2	NM3	QA1	QA2	Avg.
32K	Dense	99.67	100	100	100	89	63	91.95
32K	PQ	97.33	100	100	100	88	62	91.22
32K	vAttn (PQ)	96.33	100	100	100	89	58	90.56
32K	Oracle-top-$k$	99.33	100	100	100	88	61	91.39
32K	SSA	99.67	100	100	100	89	60	91.45
64K	Dense	99.67	100	100	100	88	59	91.11
64K	PQ	98.33	100	100	100	88	58	90.72
64K	vAttn (PQ)	99.67	100	100	100	87	57	90.61
64K	Oracle-top-$k$	99.33	100	100	100	86	57	90.39
64K	SSA	99.00	100	100	100	88	58	90.83
128K	Dense	96.67	100	100	100	84	56	89.45
128K	PQ	95.33	100	100	100	82	55	88.72
128K	vAttn (PQ)	96.67	100	100	100	82	57	89.28
128K	Oracle-top-$k$	97.33	100	100	100	82	57	89.39
128K	SSA	97.00	100	100	100	82	58	89.50

Context	Method	FWE	VT	NM2	NM3	QA1	QA2	Avg.
32K	Dense	99.67	100	100	100	93	68	93.45
32K	PQ	94.00	100	100	100	91	67	92.00
32K	vAttn (PQ)	98.00	100	100	100	89	63	91.67
32K	Oracle-top-$k$	98.33	100	99	100	91	67	92.56
32K	SSA	100.00	100	100	100	93	67	93.33
64K	Dense	99.33	100	100	100	88	66	92.22
64K	PQ	95.67	99.8	100	100	86	65	91.08
64K	vAttn (PQ)	98.00	99.8	100	99	88	63	91.30
64K	Oracle-top-$k$	98.67	100	100	100	89	66	92.28
64K	SSA	99.33	99.8	100	100	89	66	92.36
128K	Dense	99.67	100	100	100	87	61	91.28
128K	PQ	97.67	100	98	100	88	62	90.95
128K	vAttn (PQ)	98.67	100	98	100	84	63	90.61
128K	Oracle-top-$k$	99.33	100	99	100	88	62	91.39
128K	SSA	99.67	100	100	100	88	61	91.45

Context	Sparsity	Method	FWE	VT	NM2	NM3	QA1	QA2	Avg.
32K	50×	Dense	94.00	100.00	99.00	100.00	80.00	63.00	89.33
32K	50×	PQ	88.00	100.00	97.00	67.00	80.00	59.00	81.83
32K	50×	vAttn (PQ)	89.33	98.80	85.00	87.00	64.00	48.00	78.69
32K	50×	SSA	94.00	100.00	88.00	91.00	74.00	58.00	84.17
32K	50×	Oracle-top-$k$	90.67	100.00	99.00	90.00	81.00	59.00	86.61
64K	50×	Dense	81.00	99.80	100.00	100.00	75.00	61.00	86.13
64K	50×	PQ	84.33	99.80	89.00	71.00	68.00	56.00	78.02
64K	50×	vAttn (PQ)	81.33	99.80	71.00	87.00	63.00	49.00	75.19
64K	50×	SSA	81.67	99.80	80.00	95.00	69.00	54.00	79.91
64K	50×	Oracle-top-$k$	84.67	99.80	100.00	82.00	75.00	59.00	83.41
128K	50×	Dense	92.67	100.00	100.00	90.00	62.00	44.00	81.45
128K	50×	PQ	79.33	99.60	63.00	32.00	63.00	41.00	62.99
128K	50×	vAttn (PQ)	89.67	100.00	36.00	61.00	59.00	37.00	63.78
128K	50×	SSA	90.67	99.80	50.00	71.00	62.00	42.00	69.25
128K	50×	Oracle-top-$k$	85.00	100.00	93.00	67.00	64.00	46.00	75.83
32K	20×	PQ	92.33	100.00	99.00	97.00	83.00	60.00	88.56
32K	20×	vAttn (PQ)	95.00	100.00	96.00	100.00	78.00	61.00	88.33
32K	20×	SSA	93.33	100.00	97.00	99.00	77.00	63.00	88.22
32K	20×	Oracle-top-$k$	92.00	100.00	99.00	98.00	78.00	63.00	88.33
64K	20×	PQ	83.67	99.80	99.00	97.00	74.00	60.00	85.58
64K	20×	vAttn (PQ)	83.33	99.80	94.00	99.00	72.00	57.00	84.19
64K	20×	SSA	81.67	99.80	95.00	100.00	73.00	58.00	84.58
64K	20×	Oracle-top-$k$	83.00	99.80	100.00	95.00	71.00	60.00	84.80
128K	20×	PQ	85.33	100.00	88.00	80.00	66.00	44.00	77.22
128K	20×	vAttn (PQ)	91.67	99.80	77.00	89.00	62.00	43.00	77.08
128K	20×	SSA	93.33	99.80	81.00	89.00	62.00	42.00	77.86
128K	20×	Oracle-top-$k$	90.67	100.00	99.00	84.00	64.00	44.00	80.28