# AttentionIndex: Accelerating Learned Sparse Attention with a Causal Inverted-File Index > AttentionIndex replaces exhaustive key scoring in learned sparse attention with a causal inverted-file index, delivering 1.3-1.8x prefill throughput and far higher recall than HiSA at one million tokens. - **Source:** [arXiv](https://arxiv.org/abs/2611.02508) - **Published:** 2026-11-14 - **Permalink:** https://picx.dev/p/yRqpju - **Whiteboard:** https://picx.dev/p/yRqpju/image ## Summary ## 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: | Context | GLM-5.3 HiSA8K (LongBench-v2) | GLM-5.3 HiSA8K (RULER) | DeepSeek-V4-Flash HiSA4K (LongBench-v2) | DeepSeek-V4-Flash HiSA4K (RULER) | |---------|-------------------------------|-------------------------|------------------------------------------|-----------------------------------| | 16K | 83.9 | 81.7 | - | - | | 32K | 73.4 | 68.3 | 87.4 | 82.0 | | 64K | 59.3 | 55.9 | 80.3 | 64.7 | | 128K | 46.6 | 48.2 | 59.3 | 50.6 | | 256K | 41.8 | 41.4 | 46.0 | 41.8 | ### 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: | Method | Budget | Tput | Recall | Budget | Tput | Recall | |--------|--------|------|--------|--------|------|--------| | **512K tokens** | | | | | | | | Full DSA | - | 0.698 | 100.0 | - | 0.833 | 100.0 | | HiSA | 16K | 0.987 | 63.1 | 24K | 1.002 | 85.2 | | AttentionIndex | 16K | 1.000 | **86.4** | 16K | 1.000 | **89.5** | | **1M tokens** | | | | | | | | Full DSA | - | 0.548 | 100.0 | - | 0.608 | 100.0 | | HiSA | 24K | 1.048 | 69.6 | 24K | 0.971 | 76.5 | | AttentionIndex | 16K | 1.000 | **86.6** | 16K | 1.000 | **86.5** | ### 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: | Dataset | Full DSA | HiSA 16K | HiSA 24K | AttentionIndex 16K | AttentionIndex 24K | |---------|----------|----------|----------|-------------------|-------------------| | **GLM-5.3** | | | | | | | InfiniteBench | 83.0 | 75.0 | 78.6 | 80.4 | **83.3** | | RULER | 81.8 | 55.2 | 62.6 | 73.6 | **77.5** | | LongBench-v2 | 63.8 | 60.4 | 61.0 | 60.4 | **61.7** | | **DeepSeek-V4-Flash** | | | | | | | InfiniteBench | 81.16 | 80.46 | 80.19 | 81.20 | **82.03** | | RULER | 77.83 | 58.59 | 67.27 | 74.92 | **76.14** | | LongBench-v2 | 63.82 | 60.82 | 63.02 | 61.63 | **63.42** | 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. --- _Markdown view of https://picx.dev/p/yRqpju, served by PicX — AI-generated visual whiteboard summaries of research papers._