Summary (Overview)

• Key Contribution: AttentionIndex is a novel system that accelerates learned sparse attention by routing queries over an inverted-file (IVF) index of indexer keys, followed by refinement with the native indexer score, rather than exhaustively scoring every key in the prefix.
• Performance Gains: At one million tokens, AttentionIndex delivers 1.3–1.8× the prefill throughput of exhaustive selection on GLM-5.3 and DeepSeek-V4-Flash, while recovering 17.0 and 10.0 percentage points more of the exhaustive top-k compared to HiSA at matched throughput.
• Accuracy Improvements: AttentionIndex improves long-context benchmark accuracy by up to 18.4 percentage points over HiSA at the same candidate budget.
• Three Novel Mechanisms: (1) similarity-based clustering rather than contiguous-block routing, (2) causal incremental index maintenance with immediate insertion and periodic refits, and (3) learned routing representatives trained offline to better predict native indexer scores.
• Problem Identification: The paper demonstrates that key selection (not attention computation) dominates long-context prefill, taking 53–58% of prefill time at 1M tokens in production models.

---

Introduction and Theoretical Foundation

Background

Long-context inference has become a defining workload for large language models, with agents, code assistants, and document analysis routinely processing hundreds of thousands to millions of tokens per request. At these lengths, dense attention is prohibitively expensive, leading production models to adopt learned sparse attention.

Learned Sparse Attention

DeepSeek Sparse Attention (DSA) adds a lightweight learned indexer to sparse attention layers. The indexer scores all keys in the prefix, and the main attention operates only on the $k$ highest-scoring entries. The native indexer score for query token $t$ and earlier token $s$ is:

$$I_{t,s} = \sum_{j=1}^{H^I} w_{t,j}^I \operatorname{ReLU}\left(\mathbf{q}_{t,j}^I \cdot \mathbf{k}_s^I\right) \tag{1}$$

where $\mathbf{q}_{t,j}^I$ and weights $w_{t,j}^I$ of the $H^I$ indexer heads derive from token $t$, and each token contributes a single indexer key $\mathbf{k}_s^I$ shared by all heads.

The Key Selection Bottleneck

While sparse attention bounds the cost of attention computation, the indexer still scores every key. For a context of $N$ tokens, scoring and top-k selection each require $O(N)$ work per query and $O(N^2)$ work across prefill, whereas attention over selected entries requires only $O(k)$ work per query. At 1M tokens, indexer scoring and top-k selection take 53% of prefill time on DeepSeek-V4-Flash and 58% on GLM-5.3.

Limitations of Contiguous-Block Routing

Two-stage selection methods like HiSA route queries to contiguous blocks of keys ranked by mean-pooled representatives. However, contiguous token entries need not be similar in indexer space. As context grows, recall declines sharply at fixed candidate budgets:

Inverted-File (IVF) Indexing Principle

IVF indexes speed up approximate nearest-neighbor search by partitioning vectors into clusters (commonly via k-means) and restricting candidate evaluation to selected clusters. At query time, IVF ranks cluster representatives, selects a subset of clusters, and evaluates vectors within those clusters. However, three constraints complicate IVF application to learned indexers:

1. Incremental updates: New tokens must be added cheaply during prefill/decode
2. Causality: Every query must see an index built only from its own prefix
3. Non-inner-product scores: The native indexer score is not a simple inner product, so cluster means are poor proxies for member scores

---

Methodology

Design Overview

AttentionIndex replaces exhaustive selection with routing over an IVF index of indexer keys, followed by refinement with the native indexer score. The design has three main components:

1. Routing over Similarity Clusters

For each request and indexed layer, AttentionIndex maintains an IVF index over indexer keys. Each cluster stores:

• An inverted list of member key IDs
• A routing representative

A query scores every representative with the native indexer score (Equation 1) and ranks clusters. It then fills a candidate budget of $B$ keys: sink keys and unclustered open-window keys take reserved slots, with the top-ranked clusters' lists filling the rest. Refinement scores candidates exactly and passes the top-$k$ to the unchanged sparse attention kernel.

Each query scores about $C + B$ entries instead of $N$, where $C$ is the number of clusters (approximately 128 keys per cluster on GLM-5.3).

2. Causal Maintenance of a Growing Index

AttentionIndex separates immediate insertion from lazy reorganization:

• Open window: New keys enter an open window and are included directly in candidates of later queries (searchable before clustering)
• Sealing: Every $W$ keys (1,024 on GLM-5.3), the window is clustered into ~8 clusters of ~128 keys using size-capped k-means
• Refitting: Every $R$ tokens (16,384 on GLM-5.3), the entire prefix is regrouped into a new base with more clusters as the prefix grows

Causality is maintained by versioning: each seal or refit produces a new index version whose lists, representatives, and window metadata become visible together, and a query uses a version only if every key that built it lies in the query's prefix.

3. Incremental Refits

Refits reuse the previous base instead of clustering from scratch:

• Warm-start: Initialize from previous base's nonempty means plus new centers sampled from recent keys
• Bound-based acceleration: Keep for every key its assignment, an upper bound on distance to its center, and lower bounds to competing centers (Elkan, 2003; Drake & Hamerly, 2012)
• Only newly added keys and keys whose bounds no longer certify their assignment search all centers
• Distance work is proportional to $N\Delta C + MC$ instead of $NC$ (where $\Delta C$ is new centers and $M$ is keys needing full search)

4. Learned Routing Representatives

Membership uses geometric distance (cheap to update), but routing needs to predict native indexer scores. A small encoder shared by all clusters of a layer predicts a correction $\Delta r_c$ to the cluster mean $\mu_c$:

$$r_c = \mu_c + \Delta r_c$$

The encoder uses permutation-invariant summaries: mean, logarithm of cluster size, and mean/max-pooled learned features. Training minimizes cross-entropy between:

$$t_c = \frac{h_c}{\sum_j h_j} \quad \text{and} \quad p_c = \frac{n_c \exp(\alpha z_c)}{\sum_j n_j \exp(\alpha z_j)} \tag{2}$$

where $h_c$ counts reference keys in cluster $c$, $n_c$ counts causally visible members, $z_c$ is the native indexer score of $r_c$, and $\alpha$ is a training scale. The encoder learns to rank clusters by reference keys per candidate slot.

---

Empirical Validation / Results

Experimental Setup

• Models: GLM-5.3 (2,048 token keys per query) and DeepSeek-V4-Flash (512 compressed entries, compression stride of 4)
• Datasets: LongBench-v2, RULER, and InfiniteBench
• Hardware: Eight NVIDIA B200 GPUs with 8-way tensor parallelism
• Baselines: Full DSA (exhaustive selection) and HiSA (contiguous-block routing)

Indexer Retrieval Quality

AttentionIndex achieves higher mean recall than HiSA on all three datasets at every evaluated budget. At a 24K budget on LongBench-v2:

• GLM-5.3: recall improves from 67.16% to 89.70%
• DeepSeek-V4-Flash: recall improves from 75.88% to 90.66%

AttentionIndex achieves higher recall with smaller budgets: its 24K results exceed HiSA's 64K results on GLM-5.3 (62.5% fewer candidate slots).

Throughput Analysis

At 1M tokens, AttentionIndex reaches:

• 1.6–1.8× the throughput of Full DSA on GLM-5.3
• 1.3–1.6× the throughput of Full DSA on DeepSeek-V4-Flash

At comparable throughput (within 5%), AttentionIndex gains:

Decoding Latency Under Load

AttentionIndex achieves TPOT speedups of 1.40–1.45× over Full DSA, compared with 1.54–1.65× for HiSA24K.

End-to-End Accuracy

At a matched 24K budget, AttentionIndex improves accuracy over HiSA on all benchmarks:

Largest gains on RULER: 14.9 points on GLM and 8.87 points on DeepSeek-V4-Flash, recovering approximately 78% and 84% of HiSA's gap to Full DSA.

Ablation: Periodic Refitting

Every tested refit interval improves recall over sealing alone. At the serving interval of 16K:

• GLM recall rises from 0.788 to 0.883 at 24K budget (+9.49 points)
• DeepSeek-V4-Flash adds 8.29 points at 24K budget

Ablation: Learned Centroids

Learned representatives improve recall by 0.55–5.01 percentage points on GLM layers, with TTFT overhead of only +0.1% to +1.5%.

---

Theoretical and Practical Implications

Theoretical Significance

1. Key selection is the bottleneck: The paper provides empirical evidence that in production learned sparse attention, key selection rather than attention dominates long-context prefill, reframing the optimization target for sparse attention systems.

2. Similarity grouping > positional grouping: The results demonstrate that grouping keys by indexer-space similarity rather than positional contiguity substantially improves the recall-candidate-budget tradeoff, validating the IVF principle for learned indexers.

3. Causal index maintenance is feasible: The paper shows that incremental index construction can be integrated into the forward pass while preserving causality, with maintenance costs below the scoring work saved.

Practical Implications

1. Production deployment: AttentionIndex is implemented in SGLang and evaluated on production models (GLM-5.3, DeepSeek-V4-Flash), demonstrating practical viability.

2. Compatibility: The system leaves the sparse attention kernel unchanged, requiring only a swap of the candidate producer, making it easy to integrate into existing serving stacks.

3. Accuracy-recall relationship: The results show that improved indexer recall does not always produce proportional accuracy gains (e.g., LongBench-v2), suggesting model capability may be a limiting factor for certain reasoning-heavy tasks.

---

Conclusion

AttentionIndex accelerates learned sparse attention by maintaining a causal IVF index within each request and refining routed candidates with the native indexer score. Key findings:

• At a matched 24K candidate budget, AttentionIndex improves recall and task accuracy over HiSA on both GLM-5.3 and DeepSeek-V4-Flash
• Delivers 1.3–1.8× the prefill throughput of exhaustive selection at one million tokens
• Three key mechanisms: similarity-based clustering, causal incremental maintenance (sealing + refits), and learned routing representatives

Future work: Optimizing per-step decoding cost remains an open direction. The paper also notes that larger batch sizes would better amortize routing overhead and potentially close the throughput gap with HiSA during decoding.