DoSA: Dots Sparse Attention

AI summary

Overview

Research area: Efficient long-context inference for large language models — specifically, sparse attention and its token-selection ("indexer") stage on GPUs.

Technical level: Advanced. The paper assumes familiarity with attention, Top-K selection, GPU kernel design (Tensor Cores, TMEM), and the ReLU-based multi-head indexer used by DeepSeek Sparse Attention.

Scope: The paper introduces DoSA, a lossless certificate-based method for skipping full-head scoring of tokens that provably cannot enter the Top-K, together with a GPU implementation and evaluations on GLM-5.2 and DeepSeek-V3.2.

What This Paper Is About

Sparse attention only attends to a small set of keys, but in the DSA formulation, deciding which keys to attend to still requires computing a full multi-head score for every token in the context. That scoring pass becomes a dense bottleneck: at roughly one million tokens, the paper measures the indexer at 73.0% of prefill GPU time in GLM-5.2 and 90.8% in DeepSeek-V3.2. The goal is to make this selection computation conditional — computing full scores only for tokens that could actually win — while guaranteeing that the final Top-K is identical to what the original scorer would produce.

Key Contributions

1. A score certificate for the signed ReLU indexer. DoSA derives an upper bound on each token's score that is cheap to evaluate, built from a small set of retained heads, nonnegative projections onto retained anchors, joint residual allowances over omitted positive heads, and an activation-informed aggregate that lower-bounds what omitted negative heads subtract.
2. A formal exactness guarantee. The paper proves that a seed set of at least K visible keys provides a valid lower bound on the true Top-K threshold, and that rejecting any key whose upper bound falls strictly below that threshold preserves exact Top-K selection (Theorem 1, Appendix A).
3. An analysis of tightness and cost. The paper shows the chosen nonnegative projection minimizes the residual norm among single-anchor bounds of that form, that the joint residual allowance is never looser than the per-head sum, and that componentwise-majority selection of negative heads minimizes a separable bound on negative slack — then relates certificate slack to selection margins and head budgets.
4. A GPU realization. A four-stage kernel pipeline — screening dot products, bound evaluation and compaction, full-head scoring of survivors, and sparse writeback with exact selection — turns certified rejection into fewer score computations, fewer score writes, and less Top-K work, with provisions for both prefill and decode.

Main Findings

• Indexer speedups: DoSA is the fastest method in all 60 microbenchmark configurations, with 1.12×–2.74× speedups over SGLang and 1.17×–1.69× over LiteTopK on the complete Indexer.
• Speedups grow with context: At 1M tokens, gains reach 2.47×–2.74× over SGLang and 1.47×–1.69× over LiteTopK. For English inputs at 16K queries, speedup rises from 1.19× at 128K to 2.58× at 1M.
• Chunk size matters less than context length: Increasing query chunk size eightfold at 1M tokens (4K to 32K) moves speedup only from 2.47× to 2.69×, versus the much larger effect of context length.
• End-to-end TTFT: DoSA achieves 1.86× and 1.94× TTFT speedups over SGLang on GLM-5.2 and DeepSeek-V3.2 respectively, and 1.26× and 1.50× over LiteTopK, nearly halving latency at roughly 1M tokens.
• Ablation (SciFact, 16K queries, 896K keys): Screening alone cuts retrieval latency 37.4%; sparse writeback alone cuts it 23.4%; combined they cut 47.3%, a 1.90× speedup. Adding screening to writeback-only saves a further 31.1%, and adding writeback to screening-only saves 15.7%.
• The negative-head aggregate is essential: Removing it raises candidate survival from 2.65% to 95.06%, nearly eliminating all rejection and confirming that bounding the signed score is what makes the certificate tight.
• High pruning rates: Pruning exceeds 95% in 20 of 21 layers per domain, with domain averages of 97.2%–97.5%. Layer 1 is the weakest, still rejecting 89%–93%, suggesting score geometry depends more on depth than on input language or domain.
• Decoding gains are smaller but real: 1.06×, 1.11×, and 1.26× at 256K, 512K, and roughly 1M request lengths, coming from compact selection rather than reduced scoring.

Methodology in Plain English

Top-K selection does not actually need every candidate's exact score — it only needs enough evidence to rule a candidate out. DoSA exploits this in three steps.

First, score a small seed set of at least K visible tokens with the real scorer. The K-th best score in that seed set can never be higher than the true K-th best score overall, so it serves as a safe, conservative cutoff.

Second, for every other token, build a cheap upper bound on its true score and throw it away if the bound lands below the cutoff. Building that bound is the technical crux. The indexer sums ReLU activations weighted by signed head weights, so different heads need different treatment. Omitted positive heads are approximated by projecting their query vectors onto a retained head's query direction and charging the leftover difference against the key's norm — the projected part reuses a dot product already being computed. A further refund comes from noting that the residual terms all multiply the same key vector, so they can be combined into a single tighter allowance instead of summed. Omitted negative heads subtract from the score; the paper aggregates them into one query direction and uses activation frequencies measured on seed keys to decide which heads to include, since a head that is usually inactive can cancel active heads before the ReLU and weaken the bound.

Third, only surviving tokens get the full original multi-head score, and only scores above the cutoff are written out. Final selection uses original scores plus the seed records, so bounds never substitute for real values.

On the GPU, queries with similar retained-head budgets are grouped into tiles so dense Tensor Core matrix multiplies can be used for screening; keys that survive are compacted into shared streams (a key surviving at least two bounds within a group of four queries is loaded once for the whole group); and surviving scores are written as compact monotone codes with exact key IDs to cut memory traffic.

Why This Matters

The paper reframes long-context attention efficiency: raising model capacity while holding active computation roughly fixed, as mixture-of-experts does for feed-forward layers, has no clean analogue in attention if every token must still be scored. DoSA shows the selection stage can itself be made conditional without changing model outputs.

Real-world applications:

• Long-document analysis and question answering over hundreds of thousands to a million tokens, where time-to-first-token dominates user-perceived latency.
• Repository-scale code review and codebase search, one of the evaluated workload domains.
• Retrieval-augmented generation pipelines that stuff very large contexts and are latency-sensitive in serving.
• Agentic or document-processing services that must serve many concurrent long-context requests under fixed GPU budgets.

Industry relevance: The gains compound with context length, which is exactly the regime frontier serving systems are moving into. Because the method preserves the original scorer's Top-K rather than retraining or approximating the indexer, it can be dropped into existing DSA-style serving stacks (the baselines are SGLang and LiteTopK) without touching model weights.

Future Directions

• Training-time regularization of certificate slack. The authors propose encouraging an indexer whose computation is intrinsically easier to skip safely, rather than only bounding it post hoc.
• Block-level bounds. Following upper-bound pruning in document retrieval (MaxScore, WAND, Block-Max WAND), extending per-key certificates to key-block summaries could skip whole blocks without reading individual key vectors — reducing both scoring and key-access cost. The paper notes that naive block bounds are looser per candidate, which is precisely why per-key bounds were needed here.
• Closing the decode gap. Decode gains (1.06×–1.26×) trail prefill gains substantially; whether certificate reuse across decoding steps can be strengthened remains open.
• Short-context and loose-bound regimes. DoSA still scans every key, so short contexts or loose bounds can leave insufficient savings to offset screening cost; how to make screening worthwhile at moderate context lengths is unresolved.

Target Audience

Researchers and systems engineers working on long-context LLM inference, sparse attention, or GPU kernel optimization. It will be most valuable to readers already familiar with DSA-style indexers and Top-K retrieval pruning, who want to understand how classical document-retrieval upper-bound pruning translates to signed, ReLU-activated multi-head scores — and to practitioners deploying long-context serving who need evidence that the approach pays off end-to-end. Readers seeking a beginner introduction to sparse attention should look elsewhere first; the certificate construction and kernel section assume substantial background.

Authors’ abstract

Sparse attention visits few keys, yet finding them can still require scoring every key with every indexer head. We introduce DoSA (Dots Sparse Attention), a lossless method for making this computation conditional while preserving exact token selection. Its central observation is that Top-K selection requires sufficient evidence to reject a candidate, not necessarily its full score. A seed set supplies a lower bound on the winning threshold. Nonnegative head projections and joint residual bounds upper-bound positive contributions, while activation-informed aggregation bounds the magnitude subtracted by negative heads. The certificate directs full scoring toward candidates that remain eligible; exact-score filtering further reduces writeback and selection work. We establish the certificate and its exact-selection conditions, and analyze how geometry and selection margins govern its effectiveness. Our GPU implementation accelerates the complete Indexer by up to 2.74 times over SGLang and 1.69 times over LiteTopK across English, Chinese, and code workloads. End-to-end evaluation on GLM-5.2 and DeepSeek-V3.2 achieves up to 1.94 times and 1.50 times time-to-first-token speedups over the two baselines, respectively.