AttentionIndex Architecture: A Causal Inverted-File Index for Learned Sparse Attention and Taking Back Half of Million-Token Prefill

INFERENCE SYSTEMS

Production sparse attention spends more than half of million-token prefill just choosing which keys to attend to. AttentionIndex puts the selector itself behind an approximate nearest-neighbor index.

By Ossam Hashim

KEY FINDINGS

• In DeepSeek-style learned sparse attention (DSA), a lightweight indexer chooses the top-k keys for each query, but the indexer still scores every key in the prefix. At one million tokens, indexer scoring plus top-k selection takes 53-58% of prefill time on DeepSeek-V4-Flash and GLM-5.3 - more than the rest of the model combined, even with key compression and cross-layer index reuse.
• Two-stage selection over contiguous blocks (HiSA) cuts that cost but its recall decays as context grows: with an 8K-candidate budget on GLM-5.3, recall of the exhaustive top-k falls from 83.9% to 41.8% between 16K and 256K tokens on LongBench-v2, with a similar decline on DeepSeek-V4-Flash.
• AttentionIndex routes each query over an inverted-file (IVF) index that groups indexer keys by similarity rather than position, then refines the candidates with the native indexer score. The sparse attention kernel itself is unchanged; the approximation lives only in which keys reach the refine stage.
• Three production constraints are handled directly: keys become searchable immediately after they are generated, each query only sees clusters built from its own prefix (causality is enforced by index versioning, not masking), and an encoder trained offline adjusts cluster representatives so they predict the native indexer score, which is not a plain inner product.
• Maintenance stays cheap by design: new keys are clustered in small batches, and periodic refits regroup the prefix while reusing the previous clustering - only newly added keys and keys whose bounds no longer certify their assignment search all centers.
• On GLM-5.3 and DeepSeek-V4-Flash at one million tokens in SGLang, AttentionIndex delivers 1.3-1.8x the prefill throughput of exhaustive selection; at throughput within 5% of HiSA it recovers 17.0 and 10.0 percentage points more of the exhaustive top-k respectively, and it improves long-context benchmark accuracy by up to 18.4 points over HiSA at the same candidate budget.

The Selector Became the Bottleneck

Learned sparse attention moved the scaling problem, it did not remove it. DeepSeek Sparse Attention, used by DeepSeek-V3.2, DeepSeek-V4, and GLM-5.3, adds a small learned indexer next to the main attention: the indexer scores the keys in the prefix, and the expensive attention runs only over the k highest scorers. The main attention is now bounded per query, but the indexer is not - its cost still grows linearly with context. At one million tokens the balance inverts: scoring and selection consume 53-58% of prefill time on DeepSeek-V4-Flash and GLM-5.3, exceeding everything else in the model combined. Choosing keys, not attending to them, is the new bottleneck.

Why Positional Blocks Lose Recall

The natural fix is two-stage selection: route the query to a small candidate set, then run the native indexer only on those candidates. HiSA routes to contiguous blocks of keys ranked by mean-pooled representatives. The flaw is structural: contiguous tokens are not necessarily similar in indexer space or equally relevant to a query. A block holding a few high-value keys can rank low on its average and be discarded entirely. As context grows, the chance that relevant keys land in discarded blocks grows with it - hence the measured recall collapse from 83.9% to 41.8% on GLM-5.3 between 16K and 256K tokens. The authors' observation is that the native indexer already performs exhaustive search inside each selected block, so the only approximation is which keys the selected blocks contain. Grouping by similarity instead of position concentrates promising keys into fewer, denser candidate groups. That is exactly the inverted-file idea from approximate nearest-neighbor search, the same ANN toolbox this blog covered in the LEANN analysis.

An Inverted File With Causal Versions

Applying IVF inside a learned indexer runs into three constraints, and the design is mostly the answers to them. First, prefill and decode keep producing keys, so index maintenance must cost less than the scoring it saves: new keys become searchable immediately and are clustered in small batches, while periodic refits regroup the whole prefix - starting from the previous clustering rather than from scratch, and testing only newly added keys plus keys whose bounds no longer certify their old assignment. Second, causality: every query, including the centroid routing, may only see its own prefix. AttentionIndex produces a new index version at each refit - lists, representatives, and window metadata published together - and a query uses a version only if every key that built it lies inside that query's prefix. Causality becomes a versioning property, not a runtime mask. Third, the score: the native indexer is a learned function, not an inner product, so a plain centroid mean can be a poor proxy for how the indexer will score a cluster's members. An encoder trained offline and shared across the clusters of each layer adjusts the representatives to predict the native score.

What the Numbers Say

Implemented in SGLang and evaluated on GLM-5.3 and DeepSeek-V4-Flash at up to one million tokens, AttentionIndex reaches 1.3-1.8x the prefill throughput of exhaustive selection. At throughput matched within 5% of HiSA, it recovers 17.0 percentage points more of the exhaustive top-k on GLM-5.3 and 10.0 points on DeepSeek-V4-Flash. Long-context benchmark accuracy improves by up to 18.4 points over HiSA at the same candidate budget. The pattern across budgets is consistent: recall stops decaying with context because candidates are chosen by what keys are, not where they sit.

Assumptions and limitations

The worked example is illustrative, not the paper's system: real indexer keys are high-dimensional and learned, the native score is not a dot product, and cluster representatives come from an offline encoder rather than known labels. All performance figures (the 53-58% prefill share, 1.3-1.8x throughput, and the recall and accuracy deltas) are the authors' reported results on two models in one framework, SGLang; this analysis did not independently re-run them, and generalization to other models, serving stacks, and hardware is unverified. The recall measurements for HiSA and AttentionIndex depend on candidate budgets and context lengths; the reported collapse and recovery hold at the stated budgets, not universally. Refit and batching parameters trade maintenance cost against index freshness, and the paper's operating points may not be optimal elsewhere. AttentionIndex does not change the sparse attention kernel or the k budget of the main attention, and decode-side behavior is outside the prefill speedup claim.

Worked Example: Blocks Lose, Clusters Win

This toy makes the failure mode concrete. Ninety-six indexer keys are generated in eight tight similarity groups of twelve, then shuffled so that group members are scattered across positions - the realistic case, since related content rarely stays contiguous. A query similar to group zero is issued. Block routing ranks eight positional blocks of twelve by their mean representative and keeps the best two blocks (24 candidates). IVF routing ranks the eight similarity centers and keeps their members (also 24 candidates). The cluster labels stand in for the paper's offline-learned representatives, and a plain dot product stands in for the native indexer score.

import random
random.seed(3)
D, G, NG = 6, 8, 12
n = random.gauss
def vec():
    v = []
    for _ in range(D):
        v.append(n(0, 1))
    return v
def dot(a, b):
    s = 0.0
    for x, y in zip(a, b):
        s += x * y
    return s
centers = []
for _ in range(G):
    centers.append(vec())
keys = []
owner = []
for ci, c in enumerate(centers):
    for _ in range(NG):
        k = []
        for x in c:
            e = 0.3 * n(0, 1)
            k.append(x + e)
        keys.append(k)
        owner.append(ci)
both = list(zip(keys, owner))
random.shuffle(both)
keys = [b[0] for b in both]
owner = [b[1] for b in both]
q = []
for x in centers[0]:
    q.append(x + 0.2 * n(0, 1))
sc = [0.0] * len(keys)
for i, k in enumerate(keys):
    sc[i] = dot(q, k)
order = sorted(range(len(keys)),
    key=lambda i: -sc[i])
true = set(order[:8])
B = len(keys) // G
reps = []
for b in range(G):
    blk = keys[b * B:(b + 1) * B]
    s = [0.0] * D
    for k in blk:
        for j in range(D):
            s[j] += k[j]
    reps.append([x / B for x in s])
rs = [0.0] * G
for b in range(G):
    rs[b] = dot(q, reps[b])
bo = sorted(range(G),
    key=lambda b: -rs[b])
cand = set()
for b in bo[:2]:
    lo = b * B
    for i in range(lo, lo + B):
        cand.add(i)
cs = [0.0] * G
for c in range(G):
    cs[c] = dot(q, centers[c])
co = sorted(range(G),
    key=lambda c: -cs[c])
cand2 = set()
for i in range(len(keys)):
    if owner[i] in co[:2]:
        cand2.add(i)
rb = len(true & cand) / 8
ri = len(true & cand2) / 8
print("block recall:", rb)
print("ivf recall:", ri)

Sources & Provenance

• AttentionIndex: Accelerating Learned Sparse Attention with a Causal Inverted-File Index (Park and Hwu, NVIDIA, 2 Nov 2026) - arXiv abstract
• Full paper (HTML), including the GLM-5.3 and DeepSeek-V4-Flash evaluations and the index-versioning design - arXiv full text
• SGLang, the serving framework used for the implementation - SGLang on GitHub
• Prior GlobalTechDigest coverage of ANN indexing - LEANN architecture analysis

All performance numbers are attributed to the preprint as published. Sources accessed 4 November 2026. This analysis is independent editorial content for GlobalTechDigest.