DoSA Makes Sparse Attention Selection Conditional Without Approximation

Introduction

Sparse attention is designed to reduce the cost of long-context inference by visiting only a small subset of keys. In practice, however, the search for that subset can remain expensive. An indexer may still need to evaluate every candidate key across several indexer heads before it can determine the final Top-K set. When the candidate pool grows, this discovery stage can offset part of the benefit of sparse access.

DoSA, or Dots Sparse Attention, addresses this bottleneck with a lossless screening strategy. Instead of fully evaluating every candidate, it tries to prove that some candidates cannot possibly reach the current Top-K boundary. Those candidates are discarded early, while only candidates that remain eligible receive full scoring.

How the method works

• Top-K does not require every exact score: To reject a candidate, the system only needs an upper bound showing that even the candidate’s best possible score is below the current selection threshold. Exact scoring is necessary only for candidates that could still change the result.
• A seed set provides an initial threshold: DoSA begins with a seed set of candidates and uses it to obtain a lower bound on the winning threshold. This gives the screening process a concrete standard against which other candidates can be tested.
• Positive and negative contributions are bounded differently: Nonnegative head projections and joint residual bounds are used to upper-bound positive contributions. For negative heads, activation-informed aggregation limits the magnitude of the amount that can be subtracted. Together, these estimates form a certificate of whether a candidate remains eligible.
• Filtering continues after exact scoring: Once the remaining candidates are fully scored, exact-score filtering reduces writeback and selection work in addition to the earlier screening step.
• The selection semantics are preserved: The paper formalizes the certificate and the conditions under which it produces the same Top-K tokens as the original computation. The target is exact selection, not an approximate retrieval substitute.

Results and broader significance

The reported GPU implementation was evaluated on English, Chinese, and code workloads. It accelerates the complete Indexer by up to 2.74× compared with SGLang and by up to 1.69× compared with LiteTopK. In end-to-end tests, time to first token improves by up to 1.94× on GLM-5.2 and 1.50× on DeepSeek-V3.2 against the respective baselines. These gains reflect more than avoided score calculations: the method also reduces candidate writeback and downstream Top-K selection work.

The central engineering idea is to move sparsity earlier in the pipeline. A system may ultimately consume only a few tokens, but conventional indexing can still perform dense-looking work while discovering them. DoSA makes that discovery conditional, using verifiable bounds to avoid computations that cannot affect the final answer.

The paper also highlights why the benefit is workload-dependent. The geometry of the representations and the margin between selected and rejected candidates affect how tight the bounds are and how many candidates can be eliminated. A large selection margin generally gives the certificate more room to reject candidates, while loose bounds or ambiguous scores may require more full evaluations.

The reported speedups are maximum results for the tested implementations, models, and workloads rather than universal guarantees. Still, DoSA presents a useful direction for inference optimization: preserve the exact semantics of sparse attention while replacing unconditional scoring with “prove it is unnecessary, then compute only when needed.” That approach could become increasingly relevant as context windows and candidate pools continue to expand.

Source: Hugging Face Daily Papers