DoSA++: One Query, One Scan for Sparse Attention

Sparse attention attends to a few hundred keys out of a million. Picking those keys is the expensive part. DeepSeek Sparse Attention routes selection through a multi-head indexer that scores every visible key with every head, so the attention is sparse while the search that feeds it stays dense. At 1M tokens on GLM-5.2, that search takes 72% of prefill GPU time.

Work on this bottleneck splits along the dimension it attacks. IndexCache reuses index decisions across layers, so the search runs less often. LiteTopK fuses scoring with selection to cut memory traffic. HISA screens pooled blocks first, shrinking where the search looks, and MISA routes a subset of indexer heads per query, shrinking how many heads it uses. What survives all of these is the per-key operation itself: a head-wise ReLU dot product, evaluated wherever token scores are still needed.

DoSA made that operation conditional: bound a key’s score from a few heads, and reject it outright when the bound falls below a threshold the seed keys already guarantee. Rejection is certified, so the Top-K never changes.

DoSA++ attacks the same bottleneck from the other side. Instead of bounding what it skips, it predicts what to keep. One dot product per key proposes a candidate set; the original multi-head scorer then ranks those candidates. The prediction is cheap because of a structural fact about the indexer: ReLU makes it piecewise linear, so once the activation pattern is fixed, every head collapses into a single vector.

The pattern is not known in advance, but it can be estimated. DoSA++ samples a few hundred keys per query, scores them with the full indexer, and measures how often each head fires among the highest-scoring ones. Those rates become the coefficients of one folded query, which carries every head’s original signed weight and its estimated activation rate near the part of the score distribution that decides selection. The broad scan then costs one dot product per key instead of $H$, and full-head scoring is reserved for the seed keys and the surviving candidates.

This separates a broad, inexpensive search from a small, expressive ranking stage. Nothing is retrained: the folded direction is rebuilt per query from the pretrained indexer parameters and whatever is currently in the key cache.

The score is piecewise linear#

For one query, let $\mathbf q_h \in \mathbb R^D$ be head $h$, $w_h$ its weight, and $\mathbf k_i$ the key at position $i$. With $[x]_+ = \max(x, 0)$, the indexer score over the visible keys $\Omega$ is

The indexer returns the $K$ highest-scoring keys. Weights can be positive or negative: a positive head raises the score when it fires, a negative head lowers it.

Write the head dot product and its gate separately, $x_{hi} = \mathbf q_h^\top \mathbf k_i$ and $A_{hi} = \mathbf 1[x_{hi} > 0]$, and the score becomes $s_i = \sum_h w_h A_{hi} x_{hi}$. Freeze the gates at some pattern $a_h$ and the sum turns linear in the key:

The $H$ dot products on the left become one on the right. The catch is that the gates depend on the key, which is what we are trying to avoid computing. So we need a gate pattern good enough to rank keys, available before the scan starts.

Estimate the gates where ranking happens#

Selection only cares about the top of the score distribution. A head that fires on low-scoring keys and stays quiet on high-scoring ones is irrelevant to the Top-K, however often it fires overall. So DoSA++ estimates gates on the keys that matter.

Per query, it samples $m_a = 512$ evenly spaced positions, scores them with the full multi-head indexer, and keeps the top $\rho = 2\%$ by score, ties included, as a set $\mathcal T$. Each head’s activation rate is its firing frequency on that set:

Ranking the samples by the original signed score means both the positive contributions and the negative penalties shape which keys end up in $\mathcal T$. A head that fires throughout the high-score region gets a rate near one, a head that never fires gets near zero, and an intermittent head lands in between.

Substituting rates for gates gives the folded query and the predicted score:

Every head contributes, with its original signed weight intact; only the gates are approximated. There is no ReLU after the dot product — the head-wise ReLUs are already represented in the rates. On the GPU the folded vector is packed into FP8 with its own scale.

The prediction error, before quantization, is exactly

Each term is a head’s weight times how far its true gate sits from the estimated rate. Calibrating on high-scoring keys shrinks the error where ordering decides the outcome, and lets it grow freely among keys that were never going to be selected.

Folding is not head selection#

Dropping heads and folding heads are different operations. A subset scorer evaluates a few heads and ignores the rest, so the omitted heads contribute nothing — which is why DoSA has to bound them to stay exact, and why reranking cannot recover a key the screen discarded.

LISA folds instead of drops. Applying $\mathrm{ReLU}(x) = x + \mathrm{ReLU}(-x)$ splits the score exactly into a linear term plus a head-wise correction; since all heads share one key, the linear term collapses into a single dot product. LISA ranks keys by that term alone and omits the correction, which is equivalent to fixing every gate at one.

DoSA++ keeps the same one-dot-product-per-key budget but sets each gate to its estimated rate $\widehat p_h$ instead of to one. A head that rarely fires near the top of the distribution is scaled down rather than counted in full. The estimate costs 512 full-head probes per query, and in exchange the folded direction adapts to the query and its current key cache.

Scan, then refine#

Candidate selection needs a cutoff. DoSA++ scores $m_r = 4096$ further evenly spaced probes with the folded query and reads off a threshold $\theta$ targeting a candidate budget, then keeps every visible key with $\widetilde s_i \ge \theta$. The surviving set is capped per query before grouped queries share it.

Alongside the scan, a seed pool of $S \ge K$ visible keys is scored with the original $H$-head indexer; its Top-K winners $\mathcal B$ enter selection directly. The candidates $\mathcal C$ that survive filtering are then scored with the original rule, Equation (1), and the output is

Two thresholds do different jobs here. The scan cutoff $\theta$ proposes candidates from predicted scores; the seed threshold discards computed original scores that already fall below the $K$-th seed score, which trims the final selection. Ranking always uses the original score, so an approximate scan cannot distort the ordering of what it proposes — it can only fail to propose something.

DoSA’s rejections are certified, so its Top-K is exactly the full scan’s. DoSA++ trades that certificate for a cheaper broad pass and loses one to two keys per query out of 2,048, measured below.

GPU execution#

For $M$ queries and $N$ keys the folded queries form an $M \times D$ matrix. Each thread block holds a query tile $\widetilde Q_i \in \mathbb R^{B_q \times D}$ and streams key tiles through a shared-memory ring buffer while Tensor Cores compute $\widetilde S_{ij} = \widetilde Q_i K_j^\top$. Accumulators never leave the chip: SIMT threads apply scales, visibility masks, and per-query cutoffs, then pack the decisions into a bitmap. The scan writes bits, not scores.

Compaction converts the bitmap into candidate IDs. Grouped queries share candidate pages and reuse gathered key tiles, so a key that several queries want is fetched once. Refinement from that point on uses the original query tensors, weights, ReLUs, and reduction; the folded query appears nowhere in it.

Dense indexing costs $\Theta(MNHD)$ dot-product work. DoSA++ spends

where $\bar C$ is the average number of non-seed candidates scored per query after sharing. Only the $N$-sized term is multiplied by the full context, and that term has lost its factor of $H$. Everything still carrying $H$ is sized by $m_a$, $S$, or $\bar C$, all of which stay fixed as the context grows. The memory path splits the same way: a compact bitmap from the scan, gathered keys in refinement, and a small set of original-score records for selection.

Results#

We evaluate GLM-5.2-FP8 under SGLang with eight-way tensor parallelism and 16K-token prefill chunks. The baselines are native DSA as SGLang implements it, LISA with its head-fused linear score, and HISA, which screens 128-token blocks hierarchically before refining inside the survivors. Runtime measurements use a 6K candidate target; quality measurements use 8K.

Prefill#

At 1M tokens, time to first token falls from 131.4 to 56.0 seconds, a 2.35× speedup. The indexer and Top-K kernels run 4.98× faster, which drops their share of prefill GPU time from 72% to 35%.

The complete-Indexer microbenchmark replays the DoSA paper’s captures: two English, two Chinese, and two code inputs, five context lengths, four chunk sizes, and all 21 indexer-producing layers. Each cell is the median of six CUDA-event repetitions over the whole indexer, including projections, RoPE, quantization, cache stores, retrieval, and selection.

At 1M tokens with 16K chunks the speedups are 7.63×, 7.51×, and 7.57× for English, Chinese, and code. Context length drives the gain far more than chunk size does: with English inputs and 16K chunks the speedup climbs from 1.14× at 128K to 7.63× at 1M, while growing the chunk eightfold at 1M moves it from 4.85× to 8.01×. Equation (7) predicts exactly this — longer contexts add keys, and keys are what the folded scan made cheap.

Retrieval fidelity#

An approximate scan is only useful if the keys it proposes are the ones the original indexer would have chosen. We measure token-set recall against native DSA: with $\mathcal T$ the reference selection and $\widehat{\mathcal T}$ the approximate one, $\operatorname{Recall} = |\widehat{\mathcal T} \cap \mathcal T| / |\mathcal T|$.

Mean recall is 99.96% on RULER and 99.90% on LongBench v2. At $K = 2048$ that is roughly 0.7 and 2.0 missed keys per query. LISA reaches 96.83% at 128K on RULER and 94.36% on LongBench v2’s long inputs, where the same arithmetic gives 65 and 116 missed keys.

Downstream scores sit within a point of each other on both benchmarks:

To compare against HISA on identical inputs, we replay captured prefill queries: four long-context inputs per model across all indexer layers, 544 queries in total, at about 1M tokens for GLM-5.2 and 512K for Dots3 Note Preview. HISA runs at both its 8K candidate budget and its 4:1 block-compression setting.

DoSA++ averages 99.50% on GLM-5.2 and 97.47% on Dots3. HISA reaches 42.28% with 8K candidates and 86.32% with 4:1 compression, and LISA reaches 70.44%, on the same GLM-5.2 queries. The gap is wider on Dots3: 97.47% against 34.70%, 79.65%, and 44.27%.

Two things explain the spread. Block screening commits to a region before looking inside it, so a single high-scoring key in an otherwise unremarkable block can be lost with the block. And LISA’s recall varies with how far the omitted nonlinear correction moves the ranking in each layer, which is exactly what the estimated rates in Equation (4) are meant to absorb.

What this shares with DoSA, and what it does not#

Both methods run a cheap broad pass, compact the survivors, score them with the original indexer, and merge with independently scored seed keys. The pipelines are the same shape, and the final ranking in both is the original multi-head score.

The broad pass is where they part. DoSA computes an upper bound and rejects only when the bound proves the key cannot win; its output is exactly the full scan’s Top-K, and its speedup is limited by how tight those bounds are. DoSA++ predicts scores from one folded direction and accepts a small recall loss; the prediction has no guarantee attached, and the gain is correspondingly larger.

Both reuse the pretrained indexer parameters without retraining, so switching between them is a serving-time decision.

Limitations#

Recall measures agreement with native DSA’s selection, not task accuracy. Neither benchmark here resolves sub-percent differences in downstream quality.

Calibration assumes the high-scoring tail of 512 probes represents the high-scoring region of the whole cache. When the relevant keys cluster where those probes do not land, the folded direction is worse and nothing in the output says so. The scan still visits every key, so short contexts cannot amortize the probes.

Future work#

Our ongoing work combines global and local folded queries. A global direction estimates activation rates across the full context; local directions estimate them separately within context segments. Each segment supplies its own high-scoring samples, allowing the local predictor to follow changes in topic, language, or document structure that a single global estimate can average away.

We are also using query sharing to make these two views complement each other. Within a group of queries that shares candidates, some queries use global prediction and others use shared local prediction. Their candidate sets are merged before refinement with the original scorer. This gives the group both a context-wide view and region-specific views without running both scans for every query. We are evaluating how to divide the group between the two predictors and allocate their candidate budgets to improve recall while preserving the speed of the folded scan.

We will release the code and a full technical report covering DoSA++ and these global–local extensions.