PaperlayerView exact revision on arXiv

Source-linked AI summary

DoSA: Dots Sparse Attention

Fanxu Meng, Taosong Fang

arXiv:2610.13990v1cs.LG

TL;DR

Sparse attention still requires dense indexer scoring to identify which keys to visit. DoSA certifies when keys can be rejected without full scoring, and reports complete-Indexer gains across workloads alongside end-to-end TTFT speedups.

  • Problem

    Sparse attention visits few keys, but identifying them can require scoring every key with every indexer head, leaving a dense scoring bottleneck.

    Source: §1 · p. 1

  • Method

    DoSA uses seed-derived threshold lower bounds and token-specific score upper bounds to reject ineligible keys, then computes original scores for survivors.

    Source: §1 · p. 2, §2.2 · p. 3

  • Results

    Across English, Chinese, and code workloads, DoSA outperforms both baselines in all 60 configurations, with up to 2.74× complete-Indexer speedup over SGLang and 1.69× over LiteTopK.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Exact token selection can require less work than computing every token’s full score, with the amount safely skipped depending on query geometry.

    Source: §6 · p. 9

    DoSA still scans every key, and short contexts or loose bounds can leave insufficient savings to offset screening.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

DoSA targets the dense scoring bottleneck inside sparse attention with certified conditional computation that preserves the original token-selection rule. The paper reports gains in complete-Indexer and end-to-end evaluation.

Motivation

  • At approximately one million tokens, the indexer consumes 73.0% of prefill GPU time on GLM-5.2 and 90.8% on DeepSeek-V3.2.Its shares of logical matrix FLOPs are 47.4% and 83.2%, respectively.Source: §1 · p. 1
  • Prior selection systems still evaluate every head for every key, while MISA screens with fewer heads but cannot recover excluded winners during reranking.The paper frames its question as whether full-score computation can be avoided while certifying that screening discards no winner.Source: §1 · p. 1

Contributions

  • DoSA bounds token scores and fully scores only eligible candidates, using an activation-informed aggregate for omitted negative heads and bounds for omitted positive heads.Survivors receive original scores, and only scores meeting the seed-derived threshold are written for final selection.Source: §1 · pp. 1–2
  • Across English, Chinese, and code workloads, DoSA outperforms both baselines in all 60 configurations, reaching 2.74× over SGLang and 1.69× over LiteTopK.At approximately 1M tokens, TTFT speedups over SGLang are 1.86× on GLM-5.2 and 1.94× on DeepSeek-V3.2.Source: §1 · p. 2

2 METHOD

DoSA uses a seed-derived threshold and score upper bounds to reject keys that cannot enter Top-K, while retaining exact selection for survivors. Its bounds handle omitted positive and negative heads differently.

2.1 PROBLEM FORMULATION

  • For at least K visible keys, the goal is to recover the K highest-scoring keys while avoiding full-head scoring for keys that cannot enter the set.The indexer score applies ReLU to each head’s query-key dot product before signed weighting and reduction.Source: §2.1 · p. 2

2.2 SEED THRESHOLD AND RETAINED HEADS

  • The Kth-largest score among a seed set of at least K visible keys lower-bounds the true Top-K threshold, and a score upper bound below it certifies rejection.A poor seed weakens filtering; evaluating additional exact scores can only raise the threshold lower bound. Selected retained heads are computed directly, while omitted heads are bounded.Source: §2.2 · pp. 2–3

2.3 UPPER-BOUNDING OMITTED POSITIVE HEADS

  • For omitted positive heads, nonnegative projections reuse retained dot products and residual-norm allowances bound the remaining contribution.Joint residual bounds exploit the fact that omitted residuals share the same key, without additional key reads.Source: §2.3 · pp. 3–4

2.4 LOWER-BOUNDING OMITTED NEGATIVE CONTRIBUTIONS

  • For omitted negative heads, DoSA aggregates a subset to lower-bound the subtracted contribution; seed activation patterns guide the subset to reduce cancellation.The bound remains safe for other activation patterns, though it may be looser; the frequencies guide efficiency, not correctness.Source: §2.4 · p. 4

2.5 SCORE BOUNDS AND CANDIDATE REFINEMENT

  • The combined certificate prunes keys with U_i < τ, then computes original scores for survivors and keeps scores at or above τ for selection.With valid seed thresholds and no loss of live candidates, the procedure preserves exact Top-K selection; equality survives both filters, and bounds never replace final scores.Source: §2.5 · p. 4

3 GPU KERNEL DESIGN

The GPU design translates certified pruning into regular screening, sparse full-head refinement, and compact selection records. For decoding, it uses recent selections as threshold witnesses to reduce selection work.

Prefill screening

  • For prefill, query rows with similar retained-head budgets are grouped into fixed-width tiles for Tensor Core screening, while key loads overlap computation.Retained rows, folded weights, and auxiliary rows are prepared once and reused across key tiles; warp ballots compact survivor decisions for refinement.Source: §3 · p. 5

Shared key access in refinement

  • Refinement combines shared and per-query streams so a key surviving multiple bounds can be loaded once for full-head scoring across a query group.This trades some extra query-key evaluations for fewer key transfers, while filtering each original score against its query’s threshold.Source: §3 · p. 5

Compact writeback and exact selection

  • Threshold filtering writes only eligible refinement scores as compact records, reducing HBM traffic and selection input before exact Top-K selection.Boundary candidates are resolved using their original scores.Source: §3 · p. 5

Decoding

  • During decoding, recently selected distinct keys provide witnesses whose current-query scores yield a conservative Top-K threshold lower bound for compact selection.The selection procedure also maps the selected indices and updates the witnesses.Source: §3 · p. 5

4 EXPERIMENTS

Across long-context English, Chinese, and code workloads, DoSA accelerates complete-Indexer and end-to-end inference performance. Experiments vary context length, prefill chunk size, model, and ablation components to assess those gains.

4 EXPERIMENTS

  • Up to 2.74× complete-Indexer speedup and 1.94× TTFT speedup over SGLang are achieved by DoSA.DoSA is fastest in all 60 microbenchmark configurations; at 1M tokens, TTFT speedups over LiteTopK are 1.26× on GLM-5.2 and 1.50× on DeepSeek-V3.2.Source: §4 · p. 6, §4.2 · p. 7, §4.3 · p. 8

4.1 EXPERIMENTAL SETUP

  • Experiments use GLM-5.2 and DeepSeek-V3.2 on NVIDIA GPUs, with English, Chinese, and code contexts from 128K to approximately 1M tokens.Microbenchmarks run GLM-5.2 on one GPU; end-to-end inference uses eight-way tensor parallelism and model-specific prefill chunks.Source: §4.1 · pp. 6–7

4.2 INDEXER MICROBENCHMARKS

  • At 1M tokens, complete-Indexer speedups reach 2.47–2.74× over SGLang and 1.47–1.69× over LiteTopK.DoSA is fastest in all 60 configurations, with gains over SGLang exceeding 2× at every context length from 512K upward.Source: §4.2 · p. 7
  • Indexer speedup rises from 1.19× at 128K to 2.58× at 1M for English inputs with 16K queries, while changing query chunk size eightfold at 1M shifts speedup from 2.47× to 2.69×.The analysis attributes stronger gains to reducing work along the long key dimension; fusion alone leaves full multihead computation intact.Source: §4.2 · p. 7
  • Certificate pruning exceeds 95% in 20 of 21 layers per domain, with domain averages of 97.2%–97.5%.Layer 1 has the lowest pruning rate, rejecting 89%–93% of pairs; every other layer exceeds 95%.Source: §4.2 · pp. 7–8

4.3 END-TO-END INFERENCE PERFORMANCE

  • At approximately 1M tokens, TTFT speedups over SGLang are 1.86× on GLM-5.2 and 1.94× on DeepSeek-V3.2.Indexer time grows to 73.0% of prefill GPU time on GLM-5.2 and 90.8% on DeepSeek-V3.2; attention and feed-forward layers still contribute to latency.Source: §4.3 · p. 8

4.4 ABLATION STUDY

  • Combining screening and sparse writeback reduces retrieval latency by 47.3%, a 1.90× speedup over Base.Screening alone reduces latency by 37.4%, while sparse writeback alone reduces it by 23.4%; together they address both head-evaluation work and score traffic.Source: §4.4 · p. 8
  • Removing the negative aggregate raises candidate survival from 2.65% to 95.06%, nearly eliminating rejection.The ablation identifies the negative contribution as essential to a tight certificate for the signed Indexer score.Source: §4.4 · p. 8

5 RELATED WORK

Prior sparse-attention methods accelerate selection or reduce scoring work, but generally still evaluate every key with all DSA heads. DoSA adds per-key bounds to reject candidates before full scoring while preserving the original scorer’s Top-K.

Sparse attention with lightning indexers

  • DOSA certifies rejection under the original DSA scorer, unlike approaches that need not preserve its full scorer’s Top-K and cannot recover excluded tokens through reranking.The cited approaches include hierarchical indexing, cross-layer reuse, proxy scans, and head selection.Source: §5 · p. 9

Exact GPU selection and fused scoring

  • DOSA adds a score upper bound that rejects keys before full-head evaluation, while prior fused or streaming Top-K methods still score every key with all DSA heads.The comparison includes GVR, DeepSelect, LiteTopK, and radix selection.Source: §5 · p. 9

Score bounds in classical document retrieval

  • DOSA applies retrieval-style upper-bound pruning to signed DSA multihead scores, using tighter per-key bounds based on query-head geometry and activation structure.MaxScore, WAND, and Block-Max WAND use term-contribution bounds and a running Top-K threshold to skip document scoring; block-level bounds can be loose for individual keys.Source: §5 · p. 9

6 LIMITATIONS AND CONCLUSION

DOSA makes exact token selection possible without computing every token’s full score, with the amount of avoidable work governed by query geometry. Its benefits are strongest at long contexts, while screening costs and scope limits remain relevant.

6 LIMITATIONS AND CONCLUSION

  • Signed-head certificates, original-score completion, and sparse writeback enable conditional computation in the attention indexer, with query geometry determining how much work can be certified unnecessary.The head-budget analysis makes the cost of obtaining the certificate explicit.Source: §6 · p. 9
  • DOSA still scans every key, and short contexts or loose bounds may leave insufficient savings to offset screening; measurements establish computational benefits, not improved task accuracy.Numerical guarantees require conservative evaluation relative to the original scorer, and valid tied Top-K sets need not yield identical generated text.Source: §6 · p. 9
  • Future work could regularize certificate slack during training or extend per-key certificates to key-block bounds, potentially reducing scoring and key-access costs as context grows.The proposed block approach would use compact summaries to skip blocks whose bounds fall below the selection threshold.Source: §6 · p. 9

AI USE STATEMENT

The authors used generative AI tools during manuscript preparation and initial GPU-kernel drafting, and state that they take responsibility for the final work.

AI USE STATEMENT

  • Generative AI assisted manuscript drafting and revision and initial GPU-kernel implementations, which the authors subsequently revised; the authors take responsibility for the final manuscript and implementation.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

The reproducibility statement identifies where the paper’s methods, proofs, workloads, and profiling details appear. The certificate is established in real arithmetic for fixed queries over visible keys, with numerical conditions described separately.

REPRODUCIBILITY STATEMENT

  • Sections 2–4 and Appendices A–D cover the method, implementation, certificate proofs and analysis, numerical conditions, workloads, metrics, and profiling measurements.CPU checks verify algebraic relations and fixed-threshold selection.Source: REPRODUCIBILITY STATEMENT · p. 10

A CERTIFICATE VALIDITY AND EXACT SELECTION

A.1 VALIDITY OF THE SCORE CERTIFICATE

  • The positive-head projection bound uses nonnegative coefficients, and the proof combines ReLU subadditivity with conservative residual bounds to establish the score upper bound.Any finite nonnegative coefficient is valid when its residual is bounded conservatively.Source: §A.1 · p. 13

2. Then

  • Grouped positive residual bounds use one RMS value per key, while joint residual bounds exploit dependence among residual directions without additional key summaries.The joint bound’s query-side cost is quadratic in the number of omitted positive heads.Source: §A.1 · p. 13, §2 · pp. 13–14

A.2 FIXED-THRESHOLD ALGORITHM AND EXACT TOP-K SELECTION

A fixed seed threshold and score upper bounds let the algorithm reject candidates while preserving exact Top-K selection. The analysis characterizes how anchor geometry and activation-aware head selection affect certificate tightness.

A.2 FIXED-THRESHOLD ALGORITHM AND EXACT TOP-K SELECTION

  • A seed set fixes a threshold; keys with upper bounds below it are rejected, while survivors receive original scores before exact Top-K selection.With at least K distinct seeds, no lost live candidates, and a common tie order, the result matches full-domain scoring; without a prescribed tie order, it is a valid Top-K set.Source: §A.2 · p. 14

B.1 MINIMUM UNIFORM ALLOWANCE FOR A SINGLE ANCHOR

  • For one anchor, the smallest uniform allowance is the distance from q to the nonnegative ray generated by p.This optimizes the uniform norm allowance, not the complete pointwise bound or the retained-head set.Source: §B.1 · p. 15

B.2 CERTIFICATE SLACK AND ACTIVATION-AWARE HEAD SELECTION

  • When masks match the signs of the same real dot products, majority support minimizes the activation-frequency surrogate, and selected negative heads can further reduce it.Retaining the b largest a_h min(p_h, 1 − p_h) minimizes the remaining surrogate when all positive heads are retained; this does not optimize survivor count or latency. For a key, choosing its active omitted negative heads zeros both negative slack terms, while inactive heads can introduce cancellation.Source: §B.2 · pp. 15–16

B.3 SELECTION MARGINS AND SCREENING COST

Selection margins and certificate slack determine which keys survive and whether screening saves arithmetic. Numerical guarantees require conservative bounds relative to the original scorer and consistent scoring and tie rules.

B.3 SELECTION MARGINS AND SCREENING COST

  • Tighter bounds and stronger seeds narrow the survivor boundary, but the amount of boundary mass may still be large and tied keys survive.Survivors lie within {i : s_i ≥ τ* − ϵ − b} when δ_i ≤ ϵ.Source: §B.3 · p. 16
  • Screening saves logical dot products only when rejection is high enough to offset auxiliary heads; an added head pays off only if it avoids enough full rescoring.The cost model excludes planning, key summaries, Tensor Core padding, selection, and recovery; every key is still scanned, and screening adds work if all keys survive.Source: §B.3 · p. 16

C.1 REFERENCE SCORES AND THE FP8 CERTIFICATE

  • Numerical certificates must conservatively enclose representation and arithmetic errors; round-to-nearest alone is insufficient, and exact-selection conditions require the same original scorer and applicable tie rules.Invalid bounds or failure to establish the required dominance invalidate pruning guarantees.Source: §C.1 · p. 17

C.2 ERROR ALLOWANCES AND CONSERVATIVE KEY SUMMARIES

The numerical certificate uses conservative error allowances and key summaries to preserve safe pruning. Candidate storage and profiling scope impose additional operational and measurement boundaries.

C.2 ERROR ALLOWANCES AND CONSERVATIVE KEY SUMMARIES

  • Grouped Cauchy–Schwarz bounds provide a sufficient RMS error coefficient, with coordinate envelopes covering reconstruction, weight, and arithmetic errors.Underflow and other absolute errors enter ε_raw or ε_out depending on whether they occur before or after key scaling; nearest-rounded RMS summaries are insufficient.Source: §C.2 · pp. 17–18

C.3 CANDIDATE STORAGE AND FULL-DOMAIN RECOVERY

  • Candidate capacity cannot cap the number of keys passing the bound: if all scores equal the threshold, every key may pass, so overflow requires exact processing, extension, or full-domain recovery.Dropping overflow candidates can lose required Top-K IDs; recovery must preserve the selection and tie-order contract.Source: §C.3 · p. 18

D.1 INDEXER PROFILING METHODOLOGY

  • Figure 1 profiles indexer time within prefill using Nsight Systems intervals cross-checked with CUDA events; it excludes CPU tokenization and request transport and is distinct from TTFT.The measurements use two independent requests per model–length pair, cleared prefix caches, and TP8; DeepSeek results beyond declared context support computational scaling observations, not task-accuracy claims.Source: §D.1 · p. 18

D.2 CERTIFICATE COMPONENT ABLATION

Layer-wise retrieval is faster than both baselines in nearly all measured cells. Ablations show that negative-head aggregation is central to filtering, with its cost and effect varying across layers.

D.2 CERTIFICATE COMPONENT ABLATION

  • Removing negative aggregation raises candidate survival from 2.65% to 95.06% and increases Indexer time by 9.51×.Aggregating all negative heads yields 4.06% survival and 12.0× latency; strength-only selection and removing positive-head projection have much smaller aggregate effects, and projection is required in 0.42% of evaluated query-layer rows.Source: Figure 9 · p. 19, §D.2 · p. 19

D.3 LAYER-WISE CERTIFICATE BEHAVIOR

  • At layer 2, aggregating all negative heads reaches 3.41× relative Indexer latency, with its cost most pronounced in early layers.Removing the negative aggregate leaves nearly all candidates alive in most layers.Source: §D.3 · p. 19, Figure 10 · p. 20

D.4 LAYER-WISE RETRIEVAL SPEEDUPS

  • 251 of 252 (99.6%) measured comparisons beat SGLang and 250 of 252 (99.2%) beat LiteTopK.All 21 layers have median speedup above 1× against both baselines across input samples; gains reach 3.74× over SGLang and 2.08× over LiteTopK.Source: §D.4 · p. 20, Figure 11 · p. 21, Figure 12 · p. 21

Related papers

Loading 2610.13990v1…